Update train.py

train.py

@@ -9,7 +9,8 @@ os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
 output_dir = "/tmp/t5-finetuned"
 os.makedirs(output_dir, exist_ok=True)
 # Load dataset
-dataset = load_dataset("tatsu-lab/alpaca")  # Change if using your dataset
+dataset = load_dataset("tatsu-lab/alpaca")
+dataset["train"] = dataset["train"].select(range(5000))
 
 # Check dataset structure
 print("Dataset splits available:", dataset)
@@ -48,15 +49,15 @@ print("Dataset successfully split and tokenized.")
 #  Define training arguments
 training_args = TrainingArguments(
     output_dir="/tmp/t5-finetuned",
-    per_device_train_batch_size=1,  # ✅ Reduce to 1 (was 2)
-    per_device_eval_batch_size=1,  # ✅ Reduce to 1
-    num_train_epochs=1,  # Test run (increase for full fine-tuning)
-    gradient_accumulation_steps=4,  # ✅ Helps simulate larger batch size
-    logging_steps=50,
+    per_device_train_batch_size=1,
+    per_device_eval_batch_size=1,
+    num_train_epochs=1,  # ✅ Train for 1 epoch only
+    gradient_accumulation_steps=2,  # ✅ Reduce steps to speed up
+    logging_steps=100,  # ✅ Log less frequently
+    save_steps=500,  # ✅ Save less frequently
     evaluation_strategy="epoch",
     save_strategy="epoch",
     push_to_hub=False,
-    fp16=False
 )
 
 # Set up Trainer