Update train.py

train.py

@@ -9,36 +9,47 @@ from transformers import (
 
 from config import MODEL_NAME, MAX_LENGTH, DATASET_EN_ES
 
-# Load model
+# Load tokenizer + model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
 # Load dataset
 dataset = load_dataset(DATASET_EN_ES)
 
-# Preprocess function
-def preprocess(batch):
-    inputs = tokenizer(
-        batch["term"]["en"],
-        truncation=True,
-        max_length=MAX_LENGTH
+# -----------------------------
+# FIX: proper preprocessing
+# -----------------------------
+def preprocess(example):
+    source = example["term"]["en"]
+    target = example["term"]["es"]
+
+    model_inputs = tokenizer(
+        source,
+        max_length=MAX_LENGTH,
+        truncation=True
     )
 
-    targets = tokenizer(
-        batch["term"]["es"],
-        truncation=True,
-        max_length=MAX_LENGTH
+    # IMPORTANT FIX: use text_target (correct way for seq2seq)
+    labels = tokenizer(
+        text_target=target,
+        max_length=MAX_LENGTH,
+        truncation=True
     )
 
-    inputs["labels"] = targets["input_ids"]
-    return inputs
+    model_inputs["labels"] = labels["input_ids"]
+    return model_inputs
 
-dataset = dataset.map(preprocess)
+# Apply preprocessing
+tokenized_dataset = dataset.map(preprocess, remove_columns=dataset["train"].column_names)
 
+# -----------------------------
 # Data collator
+# -----------------------------
 data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
 
-# Training settings
+# -----------------------------
+# Training arguments
+# -----------------------------
 training_args = Seq2SeqTrainingArguments(
     output_dir="./my-translation-model",
     learning_rate=2e-5,
@@ -46,14 +57,17 @@ training_args = Seq2SeqTrainingArguments(
     num_train_epochs=3,
     save_strategy="epoch",
     logging_steps=50,
-    evaluation_strategy="no"
+    evaluation_strategy="no",
+    fp16=True  # faster if GPU supports it
 )
 
+# -----------------------------
 # Trainer
+# -----------------------------
 trainer = Seq2SeqTrainer(
     model=model,
     args=training_args,
-    train_dataset=dataset["train"],
+    train_dataset=tokenized_dataset["train"],
     tokenizer=tokenizer,
     data_collator=data_collator
 )