Update train.py

train.py

@@ -11,6 +11,23 @@ os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
 # Load dataset (Replace this with your dataset)
 dataset = load_dataset("tatsu-lab/alpaca")  # Example alternative dataset
 
+# Check available dataset splits
+print("Dataset splits available:", dataset)
+
+# If "test" split is missing, use a portion of "train" split
+if "test" not in dataset:
+    dataset = dataset["train"].train_test_split(test_size=0.1)  # Split 10% for testing
+
+# Tokenize dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Assign datasets
+train_dataset = tokenized_datasets["train"]
+eval_dataset = tokenized_datasets["test"]  # This is now safely created
+
+# Debug output
+print("Dataset successfully split into train and test sets")
+
 # Load model and tokenizer
 model_name = "t5-large"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
@@ -51,10 +68,11 @@ training_args = TrainingArguments(
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_datasets["train"],
-    eval_dataset=tokenized_datasets["test"],
+    train_dataset=train_dataset,  # Now correctly assigned
+    eval_dataset=eval_dataset,  # No more KeyError
 )
 
+
 # Train the model
 trainer.train()