tokenize fn

scripts/train.py

@@ -39,7 +39,6 @@ def main():
     def tokenize_function(examples):
         texts = []
         for message_list in examples["messages"]:
-            # Format the conversation history into a single string
             formatted = ""
             for msg in message_list:
                 role = msg["role"]
@@ -51,13 +50,16 @@ def main():
                 else:
                     formatted += f"{role.capitalize()}: {content}\n"
             texts.append(formatted)
-    
-        return tokenizer(
+
+        tokenized = tokenizer(
             texts,
             truncation=True,
             max_length=4096,
             padding="max_length"
         )
+        # Add labels for language modeling
+        tokenized["labels"] = tokenized["input_ids"].copy()
+        return tokenized
 
     
     tokenized_dataset = dataset.map(