Create training.py

training.py

@@ -0,0 +1,98 @@
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling
+)
+from peft import LoraConfig, get_peft_model, merge_lora_weights
+from huggingface_hub import login
+import os
+
+# ====== HF Login (optional if pushing to Hub) ======
+hf_token = os.getenv("HF_TOKEN")  # set this as env var or hardcode
+if hf_token:
+    login(token=hf_token)
+
+# ====== 1. Config ======
+model_name = "mistralai/Mistral-7B-v0.1"
+dataset_path = "tech_domains.jsonl"  # local file or HF dataset
+output_dir = "./TechChat"
+max_seq_length = 512
+
+# ====== 2. Load Dataset ======
+dataset = load_dataset("json", data_files=dataset_path)
+
+# ====== 3. Tokenizer ======
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+
+def tokenize(example):
+    text = example["instruction"] + "\n" + example["output"]
+    tokens = tokenizer(
+        text,
+        truncation=True,
+        padding="max_length",
+        max_length=max_seq_length
+    )
+    tokens["labels"] = tokens["input_ids"].copy()
+    return tokens
+
+dataset = dataset.map(tokenize, batched=True, remove_columns=dataset["train"].column_names)
+
+# ====== 4. Load Base Model ======
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# ====== 5. Apply LoRA ======
+lora_config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    target_modules=["q_proj", "v_proj"],
+    lora_dropout=0.1,
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+model = get_peft_model(model, lora_config)
+
+# ====== 6. Data Collator ======
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# ====== 7. Training Args ======
+args = TrainingArguments(
+    output_dir="./lora_tmp",
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,
+    warmup_steps=50,
+    max_steps=1000,
+    learning_rate=2e-4,
+    fp16=True,
+    logging_steps=10,
+    save_strategy="no"  # We'll save after merging
+)
+
+# ====== 8. Trainer ======
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=dataset["train"],
+    data_collator=data_collator
+)
+
+# ====== 9. Train ======
+trainer.train()
+
+# ====== 10. Merge LoRA into full model ======
+print("Merging LoRA weights into the base model...")
+model = merge_lora_weights(model)
+
+# ====== 11. Save Full Model ======
+model.save_pretrained(output_dir)
+tokenizer.save_pretrained(output_dir)
+
+print(f"✅ Full model saved at {output_dir}")
+
+# ====== 12. (Optional) Push to Hugging Face Hub ======
+# Uncomment to push
+# model.push_to_hub("hari7261/TechChat", use_temp_dir=False)
+# tokenizer.push_to_hub("hari7261/TechChat", use_temp_dir=False)