Added chat template and attention mask

added_tokens.json

@@ -1,3 +1,3 @@
 {
-  "<s>": 0
+  "<|rwkv_tokenizer_end_of_text|>": 0
 }

config.json

@@ -4,6 +4,7 @@
   ],
   "attention_hidden_size": 768,
   "auto_map": {
+    "AutoModel": "modeling_rwkv7.Rwkv7Model",
     "AutoConfig": "configuration_rwkv7.Rwkv7Config",
     "AutoModelForCausalLM": "modeling_rwkv7.Rwkv7ForCausalLM"
   },

generation_config.json

@@ -1,12 +1,13 @@
 {
     "chat_format": "chatml",
+    "bos_token_id": 0,
     "eos_token_id": 0,
     "pad_token_id": 0,
-    "max_window_size": 4096,
+    "max_window_size": 2147483647,
     "max_new_tokens": 4096,
     "do_sample": true,
-    "top_k": 0,
-    "top_p": 0.1,
-    "repetition_penalty": 1.0,
+    "top_k": 65536,
+    "top_p": 1.0,
+    "temperature": 1.0,
     "transformers_version": "4.31.1"
 }

hf_rwkv_tokenizer.py

@@ -145,7 +145,7 @@ class Rwkv6Tokenizer(PreTrainedTokenizer):
     model_input_names = ["input_ids", "attention_mask"]
 
     def __init__(
-        self, vocab_file, bos_token="<s>", eos_token="<s>", unk_token="<s>", **kwargs
+        self, vocab_file, bos_token="<|rwkv_tokenizer_end_of_text|>", eos_token="<|rwkv_tokenizer_end_of_text|>", unk_token="<|rwkv_tokenizer_end_of_text|>", **kwargs
     ):
         if not os.path.isfile(vocab_file):
             raise ValueError(

modeling_rwkv7.py

@@ -317,7 +317,7 @@ class Rwkv7SelfAttention(nn.Module):
         self.ln_x = nn.GroupNorm(H, C, eps=self.head_size * 1e-5)
 
 
-    def forward(self, hidden, state=None, v_first=None, use_cache=False, seq_mode=True):        
+    def forward(self, hidden, state=None, v_first=None, use_cache=False, seq_mode=True, attention_mask=None):        
         # Mix hidden with the previous timestep to produce key, value, receptance
         if hidden.size(1) == 1 and state is not None:
             shifted = state[0][self.layer_id]
@@ -371,6 +371,8 @@ class Rwkv7SelfAttention(nn.Module):
             rwkv7_attn_triton(r, w, k, v, -kk, kk*a, self.head_size)
             
         xx = torch.nn.functional.group_norm(xx.view(B*T,H*N), num_groups=H, weight=self.ln_x.weight, bias=self.ln_x.bias, eps = self.ln_x.eps).view(B,T,H*N)
+        if attention_mask is not None:
+            xx *= attention_mask.unsqueeze(-1)
         #x = x + ((r * k * self.r_k).view(B,T,H,N).sum(dim=-1, keepdim=True) * v.view(B,T,H,N)).view(B,T,H*N)
         xx = xx + ((r.view(B,T,H,-1)*k.view(B,T,H,-1)*self.r_k).sum(dim=-1, keepdim=True) * v.view(B,T,H,-1)).view(B,T,C)
         xx = self.output(xx * g)
@@ -435,11 +437,15 @@ class Rwkv7Block(nn.Module):
         self.attention = Rwkv7SelfAttention(config, layer_id)
         self.feed_forward = Rwkv7FeedForward(config, layer_id)
 
-    def forward(self, hidden, state=None, v_first=None, use_cache=False, output_attentions=False, seq_mode=True):
-        attention, state, v_first = self.attention(self.ln1(hidden), state=state, v_first=v_first, use_cache=use_cache, seq_mode=seq_mode)
+    def forward(self, hidden, state=None, v_first=None, use_cache=False, output_attentions=False, seq_mode=True, attention_mask=None):
+        attention, state, v_first = self.attention(
+            self.ln1(hidden) if attention_mask is None else self.ln1(hidden) * attention_mask.unsqueeze(-1) , 
+            state=state, v_first=v_first, use_cache=use_cache, seq_mode=seq_mode, attention_mask=attention_mask)
         hidden = hidden + attention
 
-        feed_forward, state = self.feed_forward(self.ln2(hidden), state=state)
+        feed_forward, state = self.feed_forward(
+            self.ln2(hidden) if attention_mask is None else self.ln2(hidden) * attention_mask.unsqueeze(-1) ,
+            state=state)
         hidden = hidden + feed_forward
 
         outputs = (hidden, state, v_first)
@@ -743,13 +749,15 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
 
         seq_mode = inputs_embeds.shape[1] > 1
         hidden_states = self.pre_ln(inputs_embeds)
+        if attention_mask is not None:
+            hidden_states *= attention_mask.unsqueeze(-1)
         v_first = None
 
         all_self_attentions = () if output_attentions else None
         all_hidden_states = () if output_hidden_states else None
         for idx, block in enumerate(self.blocks):
             hidden_states, state, v_first, attentions = block(
-                hidden_states, state=state, v_first=v_first, use_cache=use_cache, output_attentions=output_attentions, seq_mode=seq_mode
+                hidden_states, state=state, v_first=v_first, use_cache=use_cache, output_attentions=output_attentions, seq_mode=seq_mode, attention_mask=attention_mask,
             )
 
             if output_hidden_states:
@@ -759,6 +767,8 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
                 all_self_attentions = all_self_attentions + (attentions,)
 
         hidden_states = self.ln_out(hidden_states)
+        if attention_mask is not None:
+            hidden_states *= attention_mask.unsqueeze(-1)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
@@ -846,6 +856,7 @@ class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            attention_mask=attention_mask,
         )
         hidden_states = outputs[0]
 

special_tokens_map.json

@@ -1,5 +1,5 @@
 {
-  "bos_token": "<s>",
-  "eos_token": "<s>",
-  "unk_token": "<s>"
+  "bos_token": "<|rwkv_tokenizer_end_of_text|>",
+  "eos_token": "<|rwkv_tokenizer_end_of_text|>",
+  "unk_token": "<|rwkv_tokenizer_end_of_text|>"
 }

tokenizer_config.json

@@ -2,7 +2,7 @@
   "add_prefix_space": false,
   "added_tokens_decoder": {
     "0": {
-      "content": "<s>",
+      "content": "<|rwkv_tokenizer_end_of_text|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -16,11 +16,12 @@
       null
     ]
   },
-  "bos_token": "<s>",
+  "bos_token": "<|rwkv_tokenizer_end_of_text|>",
   "clean_up_tokenization_spaces": false,
-  "eos_token": "<s>",
+  "eos_token": "<|rwkv_tokenizer_end_of_text|>",
   "model_max_length": 1000000000000000019884624838656,
   "tokenizer_class": "Rwkv6Tokenizer",
-  "unk_token": "<s>",
-  "use_fast": false
+  "unk_token": "<|rwkv_tokenizer_end_of_text|>",
+  "use_fast": false,
+  "chat_template": "{{ '<|rwkv_tokenizer_end_of_text|>' }}{% for message in messages %}{% if message['role'] == 'user' %}{{'User: ' + message['content'] + '\n\n'}}{% elif message['role'] == 'system' %}{{'System: ' + message['content'] + '\n\n'}}{% elif message['role'] == 'assistant' %}{{'Assistant: ' + message['content'] + '\n\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
 }