| |
| """Train the 1-3B router (LoRA) on the patient-intent plan SFT data. |
| |
| Reuses task14's unsloth + trl recipe. Base model Qwen3-1.5B (or 3B). |
| Output: checkpoints/qwen3-{base}-router-lora/final (adapters to merge later). |
| """ |
| import json |
| import os |
|
|
| os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com") |
|
|
| import torch |
| from unsloth import FastLanguageModel, is_bfloat16_supported |
| from trl import SFTTrainer, SFTConfig |
| from datasets import Dataset |
|
|
| BASE = os.environ.get("ROUTER_BASE", "Qwen/Qwen3-1.7B") |
| OUT_ROOT = os.environ.get("ROUTER_OUT", "/workspace/TASK18/checkpoints") |
| TRAIN_JSONL = "/workspace/TASK18/data/router/router_sft_train.jsonl" |
| VAL_JSONL = "/workspace/TASK18/data/router/router_sft_val.jsonl" |
|
|
|
|
| def fmt(msgs): |
| return ("<|im_start|>user\n" + msgs[0]["content"] + "<|im_end|>\n" |
| "<|im_start|>assistant\n" + msgs[1]["content"] + "<|im_end|>\n") |
|
|
|
|
| def load(p): |
| return [json.loads(l) for l in open(p, encoding="utf-8")] |
|
|
|
|
| train = load(TRAIN_JSONL) |
| evald = load(VAL_JSONL) |
| train_ds = Dataset.from_list([{"text": fmt(r["messages"])} for r in train]) |
| eval_ds = Dataset.from_list([{"text": fmt(r["messages"])} for r in evald]) |
| print(f"train={len(train)} val={len(evald)}", flush=True) |
|
|
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name=BASE, max_seq_length=2048, dtype=None, load_in_4bit=True) |
| model = FastLanguageModel.get_peft_model( |
| model, r=16, |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", |
| "gate_proj", "up_proj", "down_proj"], |
| lora_alpha=16, lora_dropout=0, bias="none", |
| use_gradient_checkpointing="unsloth", |
| random_state=42, use_rslora=True, |
| ) |
|
|
| folder = os.path.basename(BASE.replace("/", "-")) |
| out_dir = os.path.join(OUT_ROOT, f"{folder}-router-lora") |
| config = SFTConfig( |
| output_dir=out_dir, |
| per_device_train_batch_size=16, |
| gradient_accumulation_steps=2, |
| num_train_epochs=3, |
| learning_rate=2e-4, |
| lr_scheduler_type="cosine", |
| warmup_ratio=0.05, |
| bf16=is_bfloat16_supported(), |
| logging_steps=20, |
| eval_strategy="epoch", |
| eval_steps=None, |
| save_strategy="epoch", |
| save_total_limit=2, |
| max_seq_length=2048, |
| dataset_text_field="text", |
| packing=False, |
| report_to="none", |
| optim="adamw_8bit", |
| seed=42, |
| ) |
| trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=config, |
| train_dataset=train_ds, eval_dataset=eval_ds) |
| trainer.train() |
| final = os.path.join(out_dir, "final") |
| model.save_pretrained(final) |
| tokenizer.save_pretrained(final) |
| print("TRAINING DONE ->", final) |
|
|
