Upload training script for 1.5B LoRA fine-tune

train.py

@@ -0,0 +1,159 @@
+"""sakthai_lora_train_1.5b.py
+LoRA fine-tune SakThai on Qwen2.5-1.5B-Instruct with v4 curated dataset.
+Runs on HF Jobs (t4-small / L4).
+
+After training the adapter is pushed to:
+    https://huggingface.co/Nanthasit/sakthai-context-1.5b-tools
+"""
+
+import os, sys
+
+try:
+    from datasets import load_dataset
+    from transformers import (
+        AutoTokenizer,
+        TrainingArguments,
+        Trainer,
+        DataCollatorForLanguageModeling,
+    )
+    from transformers import Qwen2ForCausalLM
+    from peft import LoraConfig, get_peft_model
+except ImportError as e:
+    print(f"❌ Missing dependency: {e}")
+    sys.exit(1)
+
+# ── Config (optimised for 16GB T4) ────────────────────────────────────────────
+BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
+DATASET = "Nanthasit/sakthai-combined-v4"
+TARGET_REPO = "Nanthasit/sakthai-context-1.5b-tools"
+MERGE_REPO = "Nanthasit/sakthai-context-1.5b-merged"
+OUTPUT_DIR = "/tmp/lora-adapter"
+
+MAX_LENGTH = 768       # longer context for tool definitions
+LR = 2e-4
+EPOCHS = 4             # more epochs on cleaner v4
+BATCH_SIZE = 1         # 1.5B is bigger — 1 per device
+GRAD_ACCUM = 16        # effective batch = 16
+WARMUP_RATIO = 0.1
+WEIGHT_DECAY = 0.01
+
+PUSH_TO_HUB = "--no-push" not in sys.argv
+
+import transformers as _tf
+_TF_MAJOR = int(_tf.__version__.split(".")[0])
+print(f"📊 transformers v{_tf.__version__}")
+
+# ── 1. Dataset ───────────────────────────────────────────────────────────────
+print(f"\n📦 Loading dataset: {DATASET}")
+ds = load_dataset(DATASET, split="train")
+print(f"   {len(ds)} examples loaded")
+
+# ── 2. Tokenizer + model ─────────────────────────────────────────────────────
+print(f"\n📥 Loading base model: {BASE_MODEL}")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = Qwen2ForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype="auto",
+    device_map="auto",
+)
+
+# Enable gradient checkpointing to save memory
+model.gradient_checkpointing_enable()
+model.config.use_cache = False
+print(f"   Model loaded ({sum(p.numel() for p in model.parameters()):,} params)")
+
+# ── 3. LoRA ──────────────────────────────────────────────────────────────────
+print("\n🔧 Applying LoRA (r=16, alpha=32)")
+lora_config = LoraConfig(
+    r=16, lora_alpha=32,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    lora_dropout=0.1, bias="none", task_type="CAUSAL_LM",
+)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+# ── 4. Format ────────────────────────────────────────────────────────────────
+def format_example(ex):
+    msgs = ex.get("messages", [])
+    tools = ex.get("tools", [])
+    text = tokenizer.apply_chat_template(
+        msgs, tools=tools or None,
+        tokenize=False, add_generation_prompt=False,
+    )
+    return {"text": text}
+
+print("\n🔄 Formatting...")
+ds = ds.map(format_example)
+
+def tok_fn(examples):
+    return tokenizer(
+        examples["text"], truncation=True,
+        max_length=MAX_LENGTH, padding="max_length",
+    )
+
+ds = ds.map(tok_fn, batched=True, remove_columns=ds.column_names)
+ds = ds.train_test_split(test_size=0.1, seed=42)
+print(f"   Train: {len(ds['train'])}  |  Eval: {len(ds['test'])}")
+
+# ── 5. Training ──────────────────────────────────────────────────────────────
+print(f"\n🏋️  Training ({EPOCHS} epochs, LR={LR}, batch={BATCH_SIZE}×{GRAD_ACCUM})")
+args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LR,
+    num_train_epochs=EPOCHS,
+    warmup_ratio=WARMUP_RATIO,
+    weight_decay=WEIGHT_DECAY,
+    fp16=True,
+    logging_steps=5,
+    save_strategy="no",
+    report_to="none",
+    remove_unused_columns=False,
+    ddp_find_unused_parameters=None,
+    optim="adamw_torch",
+)
+
+kw = dict(model=model, args=args, train_dataset=ds["train"],
+          eval_dataset=ds["test"],
+          data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False))
+if _TF_MAJOR < 5:
+    kw["tokenizer"] = tokenizer
+
+trainer = Trainer(**kw)
+trainer.train()
+
+# ── 6. Save ──────────────────────────────────────��───────────────────────────
+print(f"\n💾 Saving adapter to {OUTPUT_DIR}")
+model.save_pretrained(OUTPUT_DIR)
+tokenizer.save_pretrained(OUTPUT_DIR)
+
+# ── 7. Push ──────────────────────────────────────────────────────────────────
+if PUSH_TO_HUB:
+    print(f"\n☁️  Pushing to {TARGET_REPO}...")
+    try:
+        from huggingface_hub import HfApi
+        HfApi().upload_folder(
+            repo_id=TARGET_REPO,
+            folder_path=OUTPUT_DIR,
+            repo_type="model",
+            commit_message=f"sakthai-lora-1.5b r=16 alpha=32 epoch={EPOCHS} v4-dataset",
+        )
+        print(f"✅ Adapter at https://huggingface.co/{TARGET_REPO}")
+    except Exception as e:
+        print(f"❌ Push failed: {e}")
+else:
+    print(f"\n⏭️  Push skipped. Adapter at {OUTPUT_DIR}")
+
+print(f"""
+{'='*50}
+✅ TRAINING COMPLETE
+   Base:    {BASE_MODEL}
+   Dataset: {DATASET}
+   Adapter: {TARGET_REPO}
+   Epochs:  {EPOCHS}
+{'='*50}
+""")