159 lines
6.0 KiB
Python
159 lines
6.0 KiB
Python
"""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}
|
||
""") |