add EOT at end of a chat

EleutherAI · Jun 21, 2024 · e360e24 · e360e24
1 parent a950f8b
commit e360e24
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tools/datasets/preprocess_data_with_chat_template.py b/tools/datasets/preprocess_data_with_chat_template.py
@@ -113,6 +113,9 @@ def build_chat(
             mask.extend([-100] * len(chat_tokens))
         else:
             mask.extend(chat_tokens)
+    if tokenizer.eos_token_id is not None:
+        mask.append(tokenizer.eos_token_id if mask[-1] != -100 else -100)
+        tokens.append(tokenizer.eos_token_id)
     return tokens, mask