Update train.py

train.py

@@ -1,40 +1,29 @@
-import torch
-from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer, TrainingArguments
-from datasets import load_dataset
 import os
+from datasets import load_dataset
+from transformers import T5ForConditionalGeneration, T5Tokenizer, Trainer, TrainingArguments
 
-# Set Hugging Face cache environment variables
+# Ensure Hugging Face cache directory is writable
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["HF_DATASETS_CACHE"] = "/app/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
 
-# Load dataset (Replace this with your dataset)
-dataset = load_dataset("tatsu-lab/alpaca")  # Example alternative dataset
+# Load dataset
+dataset = load_dataset("tatsu-lab/alpaca")  # Change if using your dataset
 
-# Check available dataset splits
+# Check dataset structure
 print("Dataset splits available:", dataset)
+print("Sample row:", dataset["train"][0])
 
-# If "test" split is missing, use a portion of "train" split
+# If no 'test' split exists, create one
 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")
+    dataset = dataset["train"].train_test_split(test_size=0.1)
 
-# Load model and tokenizer
+# Load tokenizer & model
 model_name = "t5-large"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Tokenization function
-# Define tokenization function before mapping
+# Define tokenization function
 def tokenize_function(examples):
     inputs = examples["input"]  # Ensure this matches dataset key
     targets = examples["output"]
@@ -45,51 +34,41 @@ def tokenize_function(examples):
     model_inputs["labels"] = labels["input_ids"]
     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
+#  Tokenize dataset
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-# Assign train and eval datasets
+#  Assign train & eval datasets
 train_dataset = tokenized_datasets["train"]
 eval_dataset = tokenized_datasets["test"]
 
-# Debug output
-print("Dataset successfully split and tokenized")
-
+print("Dataset successfully split and tokenized.")
 
-# Apply tokenization
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
-
-# Training arguments
+#  Define training arguments
 training_args = TrainingArguments(
     output_dir="./t5-finetuned",
-    per_device_train_batch_size=2,  # Smaller batch to avoid memory errors
-    per_device_eval_batch_size=2,  # Smaller eval batch
-    save_total_limit=1,  # Keep only 1 checkpoint
-    num_train_epochs=1,  # Quick test with 1 epoch
-    logging_steps=50,  # More frequent logging
-    evaluation_strategy="epoch",  # Evaluate only at the end of the epoch
-    save_strategy="epoch",  # Save only at the end of the epoch
-    push_to_hub=False  # Avoid pushing test model to Hugging Face Hub
+    per_device_train_batch_size=2,  # Lowered to avoid memory issues
+    per_device_eval_batch_size=2,
+    num_train_epochs=1,  # Test run (increase for full fine-tuning)
+    logging_steps=50,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    push_to_hub=False  # Change to True to upload the model to HF Hub
 )
 
-# Trainer setup
+# Set up Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=train_dataset,  # Now correctly assigned
-    eval_dataset=eval_dataset,  # No more KeyError
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
 )
 
-
-# Train the model
+# Start fine-tuning
 trainer.train()
 
-# Save and push model to Hugging Face Hub
-trainer.push_to_hub("your-hf-username/t5-cover-letter")
+print("Fine-tuning complete!")
+
+# Save model locally
+trainer.save_model("./t5-finetuned")
+
+print("Model saved successfully!")