Update train_model.py

train_model.py

@@ -1,47 +1,112 @@
 # Install required packages first:
-# pip install torch transformers safetensors
+# pip install torch transformers datasets accelerate safetensors
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import Dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling
+)
 
 # -----------------------------
-# 1️⃣ Load the trained model
+# 1️⃣ Create a small custom dataset
 # -----------------------------
-model_path = "./mini_gpt_safetensor"  # folder where model was saved
+print("📥 Creating small dataset for training...")
 
-print("📥 Loading model and tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+train_texts = [
+    "Hello, my name is Ankit.",
+    "I love programming in Python.",
+    "Transformers library makes NLP easy.",
+    "PyTorch is great for deep learning.",
+    "I am learning to fine-tune GPT models."
+]
+
+test_texts = [
+    "Hello, I am training a small GPT.",
+    "Deep learning is fun!",
+    "Python is my favorite programming language."
+]
+
+# Convert to Hugging Face Dataset
+train_data = Dataset.from_dict({"text": train_texts})
+test_data = Dataset.from_dict({"text": test_texts})
+
+# -----------------------------
+# 2️⃣ Load tokenizer
+# -----------------------------
+print("📝 Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 tokenizer.pad_token = tokenizer.eos_token  # GPT models don't have pad_token
 
-model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")  # load model
+# Tokenize dataset
+def tokenize(batch):
+    return tokenizer(batch['text'], truncation=True, padding='max_length', max_length=32)
+
+train_data = train_data.map(tokenize, batched=True)
+test_data = test_data.map(tokenize, batched=True)
+
+train_data.set_format('torch', columns=['input_ids', 'attention_mask'])
+test_data.set_format('torch', columns=['input_ids', 'attention_mask'])
+
+# -----------------------------
+# 3️⃣ Load model
+# -----------------------------
+print("🤖 Loading model...")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+# -----------------------------
+# 4️⃣ Data collator
+# -----------------------------
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False
+)
 
 # -----------------------------
-# 2️⃣ Generate text
+# 5️⃣ Training arguments
 # -----------------------------
-def generate_text(prompt, max_length=50):
-    # Tokenize prompt
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    input_ids = input_ids.to(model.device)
+training_args = TrainingArguments(
+    output_dir="./mini_gpt_safetensor",
+    overwrite_output_dir=True,
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
+    num_train_epochs=3,
+    save_strategy="epoch",
+    logging_steps=10,
+    learning_rate=5e-5,
+    weight_decay=0.01,
+    fp16=True if torch.cuda.is_available() else False,
+    save_total_limit=2,
+    push_to_hub=False,
+    report_to=None,
+    optim="adamw_torch",
+    save_safetensors=True  # saves in safetensors format
+)
 
-    # Generate text
-    output_ids = model.generate(
-        input_ids,
-        max_length=max_length,
-        do_sample=True,      # for randomness
-        top_k=50,            # sample from top 50 tokens
-        top_p=0.95,          # nucleus sampling
-        temperature=0.7,
-        num_return_sequences=1
-    )
+# -----------------------------
+# 6️⃣ Trainer
+# -----------------------------
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_data,
+    eval_dataset=test_data,
+    data_collator=data_collator
+)
 
-    # Decode output
-    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output_text
+# -----------------------------
+# 7️⃣ Train model
+# -----------------------------
+print("🏋️ Training model...")
+trainer.train()
 
 # -----------------------------
-# 3️⃣ Test generation
+# 8️⃣ Save model in safetensor format
 # -----------------------------
-prompt = "Hello, I am training a mini GPT model"
-generated_text = generate_text(prompt, max_length=50)
-print("\n📝 Generated text:")
-print(generated_text)
+print("💾 Saving model in safetensors format...")
+trainer.save_model("./mini_gpt_safetensor")
+
+print("✅ Training complete and model saved!")