Further refactoring

elk/extraction/extraction.py

@@ -240,16 +240,17 @@ def extract_hiddens(
  outputs.get("decoder_hidden_states") or outputs["hidden_states"]
  )
  # Throw out layers we don't care about
- # TODO: All of RWKV's hiddens are useful. Don't throw them out.
- if not cfg.model.startswith("rwkv"):
+ if not cfg.model.startswith("BlinkDL/rwkv"):
  hiddens = [hiddens[i] for i in layer_indices]
 
  if has_per_token_states := len(hiddens[-1].shape) > 1:
  # Current shape of each element: (batch_size, seq_len, hidden_size)
  if cfg.token_loc == "first":
  hiddens = [h[..., 0, :] for h in hiddens]
  elif cfg.token_loc == "last":
- hiddens = [h[..., -1, :] if len(h.shape) >= 2 else h for h in hiddens]
+ hiddens = [
+ h[..., -1, :] if len(h.shape) >= 2 else h for h in hiddens
+ ]
  elif cfg.token_loc == "mean":
  hiddens = [h.mean(dim=-2) for h in hiddens]
  else:

elk/rwkv_lm/rwkv_hf.py

@@ -82,7 +82,7 @@ def forward(
  return response
 
  @staticmethod
- def from_pretrained(pretrained_model_name_or_path):
+ def from_pretrained(pretrained_model_name_or_path, device):
  repo_weights_paths = {
  "BlinkDL/rwkv-4-pile-1b5": "RWKV-4-Pile-1B5-20220903-8040.pth",
  "BlinkDL/rwkv-4-pile-3b": "RWKV-4-Pile-3B-20221008-8023.pth",
@@ -96,7 +96,7 @@ def from_pretrained(pretrained_model_name_or_path):
 
  weights_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename=repo_weights_paths[pretrained_model_name_or_path])
  config = RWKVConfig.from_pretrained(pretrained_model_name_or_path)
- model = RWKVModel(config, weights_path, device="cuda")
+ model = RWKVModel(config, weights_path, device)
  return model
 
 

elk/utils/hf_utils.py

@@ -21,10 +21,12 @@
 _AUTOREGRESSIVE_SUFFIXES = ["ConditionalGeneration"] + _DECODER_ONLY_SUFFIXES
 
 
-def instantiate_model(model_str: str, device: torch.device, **kwargs) -> PreTrainedModel:
+def instantiate_model(
+ model_str: str, device: torch.device, **kwargs
+) -> PreTrainedModel:
  """Instantiate a model string with the appropriate `Auto` class."""
- if model_str.startswith("rwkv"):
- return RWKVModel(device)
+ if model_str.startswith("BlinkDL/rwkv"):
+ return RWKVModel.from_pretrained(model_str, device)
 
  model_cfg = AutoConfig.from_pretrained(model_str)
  archs = model_cfg.architectures
@@ -44,9 +46,8 @@ def instantiate_model(model_str: str, device: torch.device, **kwargs) -> PreTrai
 
 def instantiate_tokenizer(model_str: str, **kwargs) -> PreTrainedTokenizerBase:
  """Instantiate a tokenizer, using the fast one iff it exists."""
- if model_str.startswith("rwkv"):
+ if model_str.startswith("BlinkDL/rwkv"):
  return RWKVTokenizer()
- # return GPT2TokenizerFast(tokenizer_file="elk/rwkv_lm/20B_tokenizer.json")
 
  try:
  return AutoTokenizer.from_pretrained(model_str, use_fast=True, **kwargs)
@@ -59,8 +60,8 @@ def instantiate_tokenizer(model_str: str, **kwargs) -> PreTrainedTokenizerBase:
 
 def instantiate_config(model_str: str, **kwargs) -> PretrainedConfig:
  """Instantiate a config."""
- if model_str.startswith("rwkv"):
- return RWKVConfig()
+ if model_str.startswith("BlinkDL/rwkv"):
+ return RWKVConfig.from_pretrained(model_str)
 
  return AutoConfig.from_pretrained(model_str, **kwargs)