#!/usr/bin/env python3 """Fine-tune a small model on the simulated ERP domain — two backends, one dataset. python scripts/finetune_erp.py # offline CPU demo (default) — runs anywhere python scripts/finetune_erp.py --backend hf # real LoRA on OpenBMB MiniCPM3-4B (needs GPU) Outputs (committed/published): backend/finetune/erp_sft.jsonl instruction-tuning dataset from the ERP KB backend/finetune/erp_finetune_report.json before→after metrics + loss curve (served at /api/erp/finetune-report and shown in the UI) backend/finetune/runs// per-run snapshot The `hf` backend builds the exact PEFT/TRL SFTTrainer config for MiniCPM3-4B and (if torch+peft+trl are installed and a GPU is present) runs it; otherwise it writes the ready-to-run recipe so it can be launched on a GPU box / HF Space / Colab unchanged. """ from __future__ import annotations import argparse import json import sys import time from pathlib import Path ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT / "backend")) from app.config import get_settings # noqa: E402 from app.erp.finetune import build_dataset, run_offline_finetune # noqa: E402 FT_DIR = ROOT / "backend" / "finetune" BASE_MODEL = "openbmb/MiniCPM3-4B" def _lora_recipe(jsonl: Path) -> dict: """The production recipe: LoRA SFT of OpenBMB MiniCPM3-4B on the ERP dataset.""" return { "base_model": BASE_MODEL, "method": "LoRA (PEFT) supervised fine-tuning (TRL SFTTrainer)", "dataset": str(jsonl.relative_to(ROOT)), "prompt_template": "{instruction}\n\nERP question: {input}\nSQL:", "hyperparams": { "lora_r": 16, "lora_alpha": 32, "lora_dropout": 0.05, "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"], "learning_rate": 2e-4, "num_train_epochs": 3, "per_device_train_batch_size": 8, "gradient_accumulation_steps": 2, "max_seq_length": 1024, "bf16": True, }, "command": "python scripts/finetune_erp.py --backend hf", "requirements": ["torch", "transformers>=4.44", "peft", "trl", "accelerate", "datasets"], } def _run_hf(jsonl: Path, settings) -> dict: """Run a real LoRA SFT if the stack is present; else emit the runnable recipe.""" recipe = _lora_recipe(jsonl) try: import torch # noqa from datasets import load_dataset # noqa from peft import LoraConfig # noqa from transformers import AutoModelForCausalLM, AutoTokenizer # noqa from trl import SFTConfig, SFTTrainer # noqa except Exception as e: return {"backend": "hf", "ran": False, "reason": f"training stack unavailable ({e})", "recipe": recipe, "note": "Dataset + recipe are ready; launch on a GPU box to fine-tune MiniCPM3-4B."} import torch from datasets import load_dataset from peft import LoraConfig from transformers import AutoModelForCausalLM, AutoTokenizer from trl import SFTConfig, SFTTrainer tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( BASE_MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, device_map="auto") ds = load_dataset("json", data_files=str(jsonl), split="train") def fmt(ex): return {"text": f"{ex['instruction']}\n\nERP question: {ex['input']}\nSQL: {ex['output']}{tok.eos_token}"} ds = ds.map(fmt) out = FT_DIR / "runs" / f"hf_{time.strftime('%Y%m%dT%H%M%S')}" trainer = SFTTrainer( model=model, train_dataset=ds, peft_config=LoraConfig(**{k: recipe["hyperparams"][k] for k in ("lora_r", "lora_alpha", "lora_dropout", "target_modules")}, task_type="CAUSAL_LM"), args=SFTConfig(output_dir=str(out), num_train_epochs=3, per_device_train_batch_size=8, learning_rate=2e-4, logging_steps=10, max_seq_length=1024, bf16=torch.cuda.is_available())) res = trainer.train() trainer.save_model(str(out)) return {"backend": "hf", "ran": True, "adapter_dir": str(out), "train_loss": float(getattr(res, "training_loss", 0.0)), "recipe": recipe} def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--backend", choices=["local", "hf"], default="local") ap.add_argument("--epochs", type=int, default=400) args = ap.parse_args() settings = get_settings() FT_DIR.mkdir(parents=True, exist_ok=True) (FT_DIR / "runs").mkdir(exist_ok=True) # 1) build + write the shared instruction-tuning dataset data = build_dataset() jsonl = FT_DIR / "erp_sft.jsonl" jsonl.write_text("\n".join(json.dumps(r) for r in data) + "\n") # 2) train if args.backend == "local": result = run_offline_finetune(settings, epochs=args.epochs) result["backend"] = "local" result["dataset_jsonl"] = str(jsonl.relative_to(ROOT)) result["production_recipe"] = _lora_recipe(jsonl) else: result = _run_hf(jsonl, settings) # always include the offline metrics too, so the UI has a populated report result["offline_demo"] = run_offline_finetune(settings, epochs=args.epochs) result["base_model_for_production"] = BASE_MODEL result["generated_at"] = time.time() # 3) publish report = FT_DIR / "erp_finetune_report.json" report.write_text(json.dumps(result, indent=2)) snap = FT_DIR / "runs" / f"{args.backend}_{time.strftime('%Y%m%dT%H%M%S')}.json" snap.write_text(json.dumps(result, indent=2)) # 4) print a readout r = result if args.backend == "local" else result.get("offline_demo", {}) print("\n" + "=" * 78) print(" ERP DOMAIN FINE-TUNE (backend: %s)" % args.backend) print("=" * 78) print(f" dataset : {len(data)} examples → {jsonl.relative_to(ROOT)}") print(f" production target : {BASE_MODEL} (LoRA recipe emitted)") if r: print(f" offline trainer : {r['model']}") print(f" classes={r['n_classes']} train={r['train']} test={r['test']} params={r['trainable_params']:,}") print(f" BEFORE test-acc : {r['before_test_accuracy']*100:5.1f}%") print(f" AFTER test-acc : {r['after_test_accuracy']*100:5.1f}% (+{r['accuracy_gain']*100:.1f} pts)") print(f" routed-SQL exec : {r['routed_sql_exec_rate']*100:.1f}% final loss={r['final_loss']}") print(f" published : {report.relative_to(ROOT)}") print("=" * 78 + "\n") if __name__ == "__main__": main()