Refactored RWKV and merged from main

.vscode/launch.json

@@ -9,7 +9,7 @@
  "type": "python",
  "request": "launch",
  "module": "elk",
- "args": ["elicit", "RWKV", "imdb", "--max_examples=5"],
+ "args": ["elicit", "rwkv", "imdb", "--max_examples=5"],
  "env": {
  "CUDA_VISIBLE_DEVICES": "5",
  },

elk/rwkv_lm/rwkv_hf.py

@@ -1,7 +1,9 @@
 import os
+import torch
 from rwkv.model import RWKV
 from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, GPT2TokenizerFast, PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutput
 
 os.environ["RWKV_JIT_ON"] = '1'
 os.environ["RWKV_CUDA_ON"] = '0'
@@ -29,9 +31,14 @@ def forward(
  position_ids=None,
  head_mask=None,
  labels=None,
+ output_hidden_states=None
  ):
- _, state = self.model.forward(input_ids, None)
- return state
+ inputs = input_ids.detach().cpu()
+ token, states = self.model.forward(inputs, None)
+ mock_embedding_state = states[0].clone()
+ output_states = [mock_embedding_state] + states
+ response = CausalLMOutput(logits=token, hidden_states=output_states)
+ return response
 
  # @staticmethod
  # def from_pretrained(pretrained_model_name_or_path):

elk/utils/hf_utils.py

@@ -6,8 +6,9 @@
  PretrainedConfig,
  PreTrainedModel,
  PreTrainedTokenizerBase,
+ GPT2TokenizerFast
 )
-from ..rwkv_lm.rwkv_hf import RWKVModel, RWKVConfig, RWKVTokenizer
+from ..rwkv_lm.rwkv_hf import RWKVModel, RWKVConfig
 
 # Ordered by preference
 _DECODER_ONLY_SUFFIXES = [
@@ -42,7 +43,7 @@ def instantiate_model(model_str: str, **kwargs) -> PreTrainedModel:
 def instantiate_tokenizer(model_str: str, **kwargs) -> PreTrainedTokenizerBase:
  """Instantiate a tokenizer, using the fast one iff it exists."""
  if model_str.startswith("rwkv"):
- return RWKVTokenizer()
+ return GPT2TokenizerFast(tokenizer_file="elk/rwkv_lm/20B_tokenizer.json")
 
  try:
  return AutoTokenizer.from_pretrained(model_str, use_fast=True, **kwargs)