Upload 4 files

app.py

@@ -172,33 +172,12 @@ def train_model(epochs, batch_size, learning_rate, resume=False, progress=gr.Pro
             greater_is_better=False,
         )
         
-        # Custom data collator that handles labels properly
-        from dataclasses import dataclass
-        from typing import Any, Dict, List
-        
-        @dataclass
-        class CustomDataCollator:
-            tokenizer: Any
-            
-            def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
-                import torch
-                
-                # Use tokenizer's pad method for proper padding
-                batch = self.tokenizer.pad(
-                    features,
-                    padding=True,
-                    return_tensors="pt"
-                )
-                
-                # Create labels from input_ids
-                # Replace padding token id with -100 so it's ignored in loss
-                labels = batch["input_ids"].clone()
-                labels[labels == self.tokenizer.pad_token_id] = -100
-                batch["labels"] = labels
-                
-                return batch
-        
-        data_collator = CustomDataCollator(tokenizer=tokenizer)
+        # Use standard data collator for causal language modeling
+        # This properly handles CoDA's internal sequence modifications
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False  # Causal LM, not masked LM
+        )
         
         # Initialize trainer with custom loss
         trainer = CoDATrainer(