Upload 4 files

app.py

@@ -181,33 +181,22 @@ def train_model(epochs, batch_size, learning_rate, resume=False, progress=gr.Pro
             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)
+                import torch
                 
-                batch = {
-                    "input_ids": [],
-                    "attention_mask": [],
-                    "labels": []
-                }
+                # Use tokenizer's pad method for proper padding
+                batch = self.tokenizer.pad(
+                    features,
+                    padding=True,
+                    return_tensors="pt"
+                )
                 
-                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)
+                # 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
                 
-                # 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)
-                }
+                return batch
         
         data_collator = CustomDataCollator(tokenizer=tokenizer)