Update train.py

train.py

@@ -9,6 +9,7 @@ import torch
 from transformers import TrainingArguments
 from trl import SFTTrainer
 
+
 def main():
     # 1) Load 4-bit model + tokenizer
     model_name = "HuggingFaceTB/SmolLM2-1.7B"
@@ -23,29 +24,34 @@ def main():
         load_in_4bit=load_in_4bit,
     )
 
-    # 2) Ensure pad token is set
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    model.config.pad_token_id = tokenizer.eos_token_id
+    # 2) Configure special tokens for chat format
+    special_tokens = {
+        "bos_token": "<|im_start|>user\n",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|im_end|>",
+        "additional_special_tokens": [
+            "<|im_start|>assistant\n"
+        ]
+    }
+    tokenizer.add_special_tokens(special_tokens)
+    model.resize_token_embeddings(len(tokenizer))
+    model.config.bos_token_id = tokenizer.bos_token_id
     model.config.eos_token_id = tokenizer.eos_token_id
+    model.config.pad_token_id = tokenizer.pad_token_id
 
-    # 3) Load instruction-response dataset
-    df = pd.read_json("cleaned_instruction_response.jsonl", lines=True)
-
-    # Rebuild clean training text
-    df["text"] = df.apply(
-        lambda row: row["instruction"].strip() + "\n\n" + row["response"].strip() + tokenizer.eos_token,
-        axis=1
-    )
+    # 3) Load chat-formatted dataset
+    df = pd.read_json("chat_formatted_data.jsonl", lines=True)
+    # Confirm each sample ends with the end-of-turn token
+    assert df['text'].str.endswith("<|im_end|>").all(), "Some samples missing end-of-turn token"
 
-    # Convert to Hugging Face Dataset
+    # 4) Create Hugging Face Dataset and split
     full_dataset = Dataset.from_pandas(df[["text"]])
     split = full_dataset.train_test_split(test_size=0.15, seed=42)
     train_dataset = split["train"]
     eval_dataset = split["test"]
     print(f"✅ Training samples: {len(train_dataset)} | Eval samples: {len(eval_dataset)}")
 
-    # 4) Apply LoRA adapters
+    # 5) Apply LoRA adapters
     model = FastLanguageModel.get_peft_model(
         model,
         r=8,
@@ -59,7 +65,7 @@ def main():
         use_gradient_checkpointing=True,
     )
 
-    # 5) Tokenization
+    # 6) Tokenization function
     def tokenize_fn(examples):
         tokens = tokenizer(
             examples["text"],
@@ -73,7 +79,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) Training arguments
+    # 7) Training arguments
     training_args = TrainingArguments(
         output_dir="./output_model",
         per_device_train_batch_size=8,
@@ -92,7 +98,7 @@ def main():
         dataloader_num_workers=2,
     )
 
-    # 7) Train
+    # 8) Train with SFTTrainer
     trainer = SFTTrainer(
         model=model,
         tokenizer=tokenizer,
@@ -105,13 +111,14 @@ def main():
     trainer.train()
     print("✅ Training complete.")
 
-    # 8) Save and zip
+    # 9) Save and zip model
     final_dir = "./output_model_final"
     os.makedirs(final_dir, exist_ok=True)
     model.save_pretrained(final_dir)
     tokenizer.save_pretrained(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:
@@ -121,4 +128,4 @@ def main():
     print(f"✅ Model zipped → {zip_path}")
 
 if __name__ == "__main__":
-    main()
+    main()