Update train_model.py

train_model.py

@@ -2,8 +2,13 @@ from datasets import load_dataset
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from transformers import Trainer, TrainingArguments
 
+ds = load_dataset("openai/summarize_from_feedback", "axis")
+
 # Load dataset (replace with your actual dataset path or URL)
 dataset = load_dataset("path_to_your_dataset")
+from datasets import load_dataset
+
+
 
 # Load the pre-trained model and tokenizer
 model_name = "gpt2"  # Or another model like T5, BART, etc.