Update train.py

train.py

@@ -34,10 +34,9 @@ tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 # Tokenization function
+# Define tokenization function before mapping
 def tokenize_function(examples):
-    print("Sample data structure:", examples)  # Move print inside function
-
-    inputs = examples["input"]  # Make sure "input" matches dataset keys
+    inputs = examples["input"]  # Ensure this matches dataset key
     targets = examples["output"]
 
     model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
@@ -46,7 +45,23 @@ def tokenize_function(examples):
     model_inputs["labels"] = labels["input_ids"]
     return model_inputs
 
-    return model_inputs
+# Check dataset structure
+print("Dataset splits available:", dataset)
+
+# If "test" split is missing, create one
+if "test" not in dataset:
+    dataset = dataset["train"].train_test_split(test_size=0.1)
+
+# Tokenize dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Assign train and eval datasets
+train_dataset = tokenized_datasets["train"]
+eval_dataset = tokenized_datasets["test"]
+
+# Debug output
+print("Dataset successfully split and tokenized")
+
 
 # Apply tokenization
 tokenized_datasets = dataset.map(tokenize_function, batched=True)