Update train.py

train.py

@@ -3,8 +3,8 @@ from unsloth import FastLanguageModel
 
 import os
 import zipfile
-from datasets import Dataset
 import pandas as pd
+from datasets import Dataset
 import torch
 from transformers import TrainingArguments
 from trl import SFTTrainer
@@ -27,20 +27,17 @@ def main():
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     model.config.pad_token_id = tokenizer.eos_token_id
+    model.config.eos_token_id = tokenizer.eos_token_id
 
-    # 3) Load and clean dataset
-    df = pd.read_json("data.jsonl", lines=True)
+    # 3) Load instruction-response dataset
+    df = pd.read_json("cleaned_instruction_response.jsonl", lines=True)
 
-    # Ensure all rows end cleanly with a single eos_token
-    df["text"] = df["text"].apply(
-        lambda t: t.rstrip()
-        .removesuffix(tokenizer.eos_token)
-        .strip() + tokenizer.eos_token
+    # Rebuild clean training text
+    df["text"] = df.apply(
+        lambda row: row["instruction"].strip() + "\n\n" + row["response"].strip() + tokenizer.eos_token,
+        axis=1
     )
 
-    # Confirm all rows end properly
-    assert df["text"].str.endswith(tokenizer.eos_token).all(), "Some rows are missing eos_token!"
-
     # Convert to Hugging Face Dataset
     full_dataset = Dataset.from_pandas(df[["text"]])
     split = full_dataset.train_test_split(test_size=0.15, seed=42)
@@ -62,7 +59,7 @@ def main():
         use_gradient_checkpointing=True,
     )
 
-    # 5) Tokenize and retain eos token
+    # 5) Tokenization
     def tokenize_fn(examples):
         tokens = tokenizer(
             examples["text"],
@@ -76,7 +73,7 @@ def main():
     tokenized_train = train_dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
     tokenized_eval = eval_dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
 
-    # 6) Define training arguments
+    # 6) Training arguments
     training_args = TrainingArguments(
         output_dir="./output_model",
         per_device_train_batch_size=8,
@@ -88,15 +85,14 @@ def main():
         logging_steps=25,
         save_strategy="epoch",
         save_total_limit=2,
-        eval_strategy="epoch",  # ✅ Corrected
+        eval_strategy="epoch",
         load_best_model_at_end=True,
         metric_for_best_model="eval_loss",
         greater_is_better=False,
         dataloader_num_workers=2,
     )
 
-
-    # 7) Train the model
+    # 7) Train
     trainer = SFTTrainer(
         model=model,
         tokenizer=tokenizer,
@@ -109,22 +105,20 @@ def main():
     trainer.train()
     print("✅ Training complete.")
 
-    # 8) Save model and zip it
+    # 8) Save and zip
     final_dir = "./output_model_final"
     os.makedirs(final_dir, exist_ok=True)
     model.save_pretrained(final_dir)
     tokenizer.save_pretrained(final_dir)
-    print(f"📦 Saved final model to {final_dir}")
 
     zip_path = "model.zip"
-    print(f"🗜 Zipping model → {zip_path}")
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
         for root, _, files in os.walk(final_dir):
             for fname in files:
                 full = os.path.join(root, fname)
                 rel = os.path.relpath(full, final_dir)
                 zf.write(full, rel)
-    print(f"✅ Zipped model saved to {zip_path}")
+    print(f"✅ Model zipped → {zip_path}")
 
 if __name__ == "__main__":
-    main()
+    main()