scripts/finetune_erp.py

#!/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/<ts>/               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()