初始化项目,由ModelHub XC社区提供模型
Model: Nanthasit/sakthai-context-1.5b-merged Source: Original Platform
This commit is contained in:
159
train.py
Normal file
159
train.py
Normal file
@@ -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}
|
||||
""")
|
||||
Reference in New Issue
Block a user