Update train.py

train.py

@@ -17,15 +17,20 @@ tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 # Tokenization function
-print("Sample data structure:", examples)
 def tokenize_function(examples):
-    inputs = [ex["input"] for ex in examples]
-    targets = [ex["output"] for ex in examples]
+    print("Sample data structure:", examples)  # Move print inside function
+
+    inputs = examples["input"]  # Make sure "input" matches dataset keys
+    targets = examples["output"]
+
     model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
     labels = tokenizer(targets, max_length=512, truncation=True, padding="max_length")
+
     model_inputs["labels"] = labels["input_ids"]
     return model_inputs
 
+    return model_inputs
+
 # Apply tokenization
 tokenized_datasets = dataset.map(tokenize_function, batched=True)