Update train_model.py

train_model.py

@@ -1,42 +1,47 @@
 from datasets import load_dataset
 from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments
 
-# Load your dataset
-dataset = load_dataset("json", data_files="dataset.jsonl")["train"]
+# Step 1: Load dataset
+dataset = load_dataset("DetectiveShadow/MVPQuestion")["train"]
 
-# Load tokenizer and model
+# Optional: Rename columns if needed
+# dataset = dataset.rename_columns({"your_input_column": "input", "your_output_column": "output"})
+
+# Step 2: Load tokenizer and model
 model_name = "t5-small"
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# Preprocessing
+# Step 3: Tokenization function
 def tokenize(example):
     input_enc = tokenizer(example["input"], truncation=True, padding="max_length", max_length=64)
     target_enc = tokenizer(example["output"], truncation=True, padding="max_length", max_length=64)
     input_enc["labels"] = target_enc["input_ids"]
     return input_enc
 
-tokenized_data = dataset.map(tokenize)
+tokenized = dataset.map(tokenize)
 
-# Training arguments
+# Step 4: Training configuration
 training_args = TrainingArguments(
-    output_dir="./trivia-genie-t5",
+    output_dir="./MVPTrivia",
     per_device_train_batch_size=8,
     num_train_epochs=3,
     logging_steps=10,
-    save_total_limit=2,
-    save_strategy="epoch"
+    save_strategy="epoch",
+    push_to_hub=True,
+    hub_model_id="DetectiveShadow/MVPTrivia",  # This is where your model will go
+    hub_strategy="every_save"
 )
 
+# Step 5: Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_data,
+    train_dataset=tokenized
 )
 
-# Train
+# Step 6: Train and push
 trainer.train()
+trainer.push_to_hub()
+tokenizer.push_to_hub("DetectiveShadow/MVPTrivia")
 
-# Optional: Push to Hugging Face Hub
-# model.push_to_hub("your-username/trivia-genie-t5")
-# tokenizer.push_to_hub("your-username/trivia-genie-t5")