Update README.md

README.md

@@ -250,8 +250,11 @@ def seallm_7b_v25_tokenize_multi_turns(tokenizer, conversations, add_assistant_p
     """
     TURN_TEMPLATE = "<|im_start|>{role}\n{content}<eos>\n"
     TURN_PREFIX = "<|im_start|>{role}\n"
+    TURN_SUFFIX = "<eos>\n"
+    TURN_SUFFIX_TAKE = "<eos>"
     sample = None
     assistant_prefix_len = None
+    assistant_suffix_len = None
     for turn_id, turn in enumerate(conversations):
         prompt = TURN_TEMPLATE.format(role=turn['role'], content=turn['content'])
         turn_sample = tokenizer(
@@ -261,7 +264,12 @@ def seallm_7b_v25_tokenize_multi_turns(tokenizer, conversations, add_assistant_p
         if turn['role'] == 'assistant':
             if assistant_prefix_len is None:
                 assistant_prefix_len = len(tokenizer.encode(TURN_PREFIX.format(role=turn['role']), add_special_tokens=False))
-            turn_sample['token_type_ids'][assistant_prefix_len:] = [1] * (len(turn_sample['input_ids']) - assistant_prefix_len)
+            if assistant_suffix_len is None:
+                assistant_suffix_len = (
+                    len(tokenizer.encode(TURN_SUFFIX.format(role=turn['role']), add_special_tokens=False)) - 
+                    len(tokenizer.encode(TURN_SUFFIX_TAKE, add_special_tokens=False))
+                )
+            turn_sample['token_type_ids'][assistant_prefix_len:-assistant_suffix_len] = [1] * (len(turn_sample['input_ids']) - assistant_prefix_len - assistant_suffix_len)
         if sample is None:
             sample = turn_sample
         else:
@@ -282,9 +290,12 @@ def seallm_7b_v25_tokenize_multi_turns(tokenizer, conversations, add_assistant_p
 
 # ! testing
 sample = seallm_7b_v25_tokenize_multi_turns(tokenizer, conversations)
-print(tokenizer.convert_ids_to_tokens(sample['input_ids']))
-print(sample['token_type_ids'])
+tokens = tokenizer.convert_ids_to_tokens(sample['input_ids'])
+pairs = [(x, y) for x, y in zip(tokens, sample['token_type_ids'])]
+print(pairs)
 
+# source and special tokens is masked out (token_type 0), only assistant with <eos> is trained (token_type 1)
+# [('<bos>', 0), ('<', 0), ('|', 0), ..., ('assistant', 0), ('\n', 0), ('Hi', 1), ('▁there', 1), (',', 1), ('▁how', 1), ('▁can', 1), ('▁I', 1), ('▁help', 1), ('?', 1), ('<eos>', 1), ('\n', 0), ('<', 0), ... 
 
 ```