Upload 4 files

app.py

@@ -133,15 +133,8 @@ 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(
-            tokenize_function,
+            lambda x: preprocess_conversations(x, tokenizer),
             batched=True,
             remove_columns=dataset.column_names
         )
@@ -179,12 +172,44 @@ def train_model(epochs, batch_size, learning_rate, resume=False, progress=gr.Pro
             greater_is_better=False,
         )
         
-        # Data collator with padding
-        data_collator = DataCollatorForLanguageModeling(
-            tokenizer=tokenizer,
-            mlm=False,
-            pad_to_multiple_of=8  # Pad to multiple of 8 for efficiency
-        )
+        # 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]:
+                # Pad input_ids
+                max_length = max(len(f["input_ids"]) for f in features)
+                
+                batch = {
+                    "input_ids": [],
+                    "attention_mask": [],
+                    "labels": []
+                }
+                
+                for f in features:
+                    input_ids = f["input_ids"]
+                    padding_length = max_length - len(input_ids)
+                    
+                    # Pad input_ids and attention_mask
+                    batch["input_ids"].append(input_ids + [self.tokenizer.pad_token_id] * padding_length)
+                    batch["attention_mask"].append(f["attention_mask"] + [0] * padding_length)
+                    
+                    # Labels: copy of input_ids with padding as -100 (ignored in loss)
+                    batch["labels"].append(input_ids + [-100] * padding_length)
+                
+                # Convert to tensors
+                import torch
+                return {
+                    "input_ids": torch.tensor(batch["input_ids"], dtype=torch.long),
+                    "attention_mask": torch.tensor(batch["attention_mask"], dtype=torch.long),
+                    "labels": torch.tensor(batch["labels"], dtype=torch.long)
+                }
+        
+        data_collator = CustomDataCollator(tokenizer=tokenizer)
         
         # Initialize trainer with custom loss
         trainer = CoDATrainer(