Upload 4 files

app.py

@@ -63,10 +63,7 @@ def preprocess_conversations(examples, tokenizer):
                 text += f"<|assistant|>\n{content}\n"
         texts.append(text)
     
-    # Return tokenized data with labels for language modeling
-    tokenized = tokenizer(texts, truncation=True, max_length=2048, padding=False)
-    tokenized["labels"] = tokenized["input_ids"].copy()
-    return tokenized
+    return tokenizer(texts, truncation=True, max_length=2048, padding=False)
 
 # Persistent storage paths
 CHECKPOINT_DIR = Path("/data/checkpoints") if Path("/data").exists() else Path("./checkpoints")
@@ -135,8 +132,16 @@ def train_model(epochs, batch_size, learning_rate, resume=False, progress=gr.Pro
         
         # Preprocess dataset
         progress(0.3, desc="Preprocessing dataset...")
+        
+        def tokenize_function(examples):
+            # Process conversations
+            processed = preprocess_conversations(examples, tokenizer)
+            # Add labels (copy of input_ids for language modeling)
+            processed["labels"] = [ids[:] for ids in processed["input_ids"]]
+            return processed
+        
         tokenized_dataset = dataset.map(
-            lambda x: preprocess_conversations(x, tokenizer),
+            tokenize_function,
             batched=True,
             remove_columns=dataset.column_names
         )