training/train_adapter.py

"""
Train one LoRA adapter per mode and push to Hugging Face Hub.

Each adapter teaches the base model to emit  <ui>{json}</ui>  then prose
for its mode, so the handler in hf-endpoint/handler.py can parse it.

GPU requirements
----------------
QLoRA (default, --no-qlora not set):  ~8 GB VRAM  (RTX 3080 / T4 / L4)
bf16 full precision (--no-qlora):     ~24 GB VRAM  (A10G / A100)

Colab/Kaggle tip: use the free T4 with QLoRA for quick experiments.

Usage
-----
    # Login first
    huggingface-cli login

    # Train one mode (generates data automatically)
    python training/train_adapter.py --mode support --hub-org your-org

    # Train all three modes back-to-back
    python training/train_adapter.py --all --hub-org your-org

    # Use your own JSONL instead of the synthetic generator
    python training/train_adapter.py --mode form --data data/form.jsonl --hub-org your-org

    # Smoke-test with tiny run (no push)
    python training/train_adapter.py --mode support --hub-org your-org --n 20 --epochs 1 --dry-run
"""
import argparse
import json
import os
import sys
from pathlib import Path
from typing import Optional

import torch
from datasets import Dataset
from peft import LoraConfig, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

# Allow running from repo root or training/
sys.path.insert(0, str(Path(__file__).parent.parent))
from training.generate_examples import _GENERATORS, generate

# ── constants ─────────────────────────────────────────────────────────────────

DEFAULT_BASE = os.getenv("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")

# Covers Qwen2.5, Llama-3.x, Mistral-0.3.
# For Phi-3/3.5 replace with ["qkv_proj", "o_proj", "gate_up_proj", "down_proj"].
LORA_TARGETS = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]


# ── data helpers ──────────────────────────────────────────────────────────────

def load_examples(mode: str, n: int, jsonl_path: Optional[str]) -> list[dict]:
    if jsonl_path:
        p = Path(jsonl_path)
        if not p.exists():
            raise FileNotFoundError(f"JSONL not found: {p}")
        rows = [json.loads(l) for l in p.read_text(encoding="utf-8").splitlines() if l.strip()]
        print(f"  Loaded {len(rows)} examples from {p}")
        return rows

    rows = generate(mode, n)
    print(f"  Generated {len(rows)} synthetic examples for '{mode}'")
    return rows


def build_dataset(examples: list[dict], tokenizer) -> tuple[Dataset, Dataset]:
    texts = [
        tokenizer.apply_chat_template(
            ex["messages"], tokenize=False, add_generation_prompt=False
        )
        for ex in examples
    ]
    full = Dataset.from_dict({"text": texts})
    split = full.train_test_split(test_size=0.1, seed=42)
    return split["train"], split["test"]


# ── model helpers ─────────────────────────────────────────────────────────────

def load_base_model(base: str, use_qlora: bool) -> AutoModelForCausalLM:
    bnb_cfg = None
    if use_qlora:
        bnb_cfg = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )

    model = AutoModelForCausalLM.from_pretrained(
        base,
        quantization_config=bnb_cfg,
        torch_dtype=torch.bfloat16 if not use_qlora else None,
        device_map="auto",
        trust_remote_code=True,
    )
    model.config.use_cache = False
    return model


def lora_config() -> LoraConfig:
    return LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules=LORA_TARGETS,
        task_type=TaskType.CAUSAL_LM,
        bias="none",
    )


# ── training ──────────────────────────────────────────────────────────────────

def train_one(
    mode: str,
    base_model: str,
    hub_org: str,
    n_examples: int,
    epochs: int,
    jsonl_path: Optional[str],
    output_root: Path,
    use_qlora: bool,
    dry_run: bool,
) -> None:
    print(f"\n{'='*60}")
    print(f"  Adapter: {mode}")
    print(f"  Base:    {base_model}")
    print(f"  QLoRA:   {use_qlora}")
    print(f"{'='*60}")

    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"   # prevents warning with flash-attn

    examples       = load_examples(mode, n_examples, jsonl_path)
    train_ds, eval_ds = build_dataset(examples, tokenizer)
    print(f"  Train: {len(train_ds)}  Eval: {len(eval_ds)}")

    if dry_run:
        print("  [dry-run] skipping model load, training, and push")
        sample = train_ds[0]["text"][:200]
        print(f"  Sample text:\n  {sample!r}…\n")
        return

    model    = load_base_model(base_model, use_qlora)
    out_dir  = output_root / f"adapter-{mode}"

    sft_cfg = SFTConfig(
        output_dir=str(out_dir),
        num_train_epochs=epochs,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,     # effective batch size 8
        gradient_checkpointing=True,       # saves ~30 % VRAM
        learning_rate=2e-4,
        warmup_ratio=0.05,
        lr_scheduler_type="cosine",
        bf16=not use_qlora,
        fp16=False,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        max_seq_length=1024,
        dataset_text_field="text",
        packing=True,                      # packs short examples → less padding waste
        report_to="none",
    )

    trainer = SFTTrainer(
        model=model,
        args=sft_cfg,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        peft_config=lora_config(),
        processing_class=tokenizer,
    )

    print(f"\n  Training {mode} adapter…")
    trainer.train()

    trainer.save_model(str(out_dir))
    tokenizer.save_pretrained(str(out_dir))
    print(f"  Saved locally → {out_dir}")

    hub_repo = f"{hub_org}/adapter-{mode}"
    print(f"  Pushing → https://huggingface.co/{hub_repo}")
    trainer.push_to_hub(hub_repo, commit_message=f"train: {mode} adapter epoch {epochs}")
    tokenizer.push_to_hub(hub_repo)
    print(f"  Done: https://huggingface.co/{hub_repo}")


# ── CLI ───────────────────────────────────────────────────────────────────────

def main() -> None:
    p = argparse.ArgumentParser(
        description="Train LoRA adapters for adaptive-model and push to HF Hub",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument("--mode", choices=list(_GENERATORS.keys()),
                   help="Single adapter mode to train")
    p.add_argument("--all", dest="train_all", action="store_true",
                   help="Train all modes sequentially")
    p.add_argument("--hub-org", required=True,
                   help="HF Hub org or username (e.g. your-org)")
    p.add_argument("--base-model", default=DEFAULT_BASE,
                   help=f"Base model ID on HF Hub (default: {DEFAULT_BASE})")
    p.add_argument("--n", type=int, default=200,
                   help="Synthetic examples to generate per mode (default: 200)")
    p.add_argument("--epochs", type=int, default=3,
                   help="Training epochs per adapter (default: 3)")
    p.add_argument("--data", metavar="JSONL",
                   help="Path to existing JSONL — skips generation (single --mode only)")
    p.add_argument("--output-dir", default="./trained-adapters",
                   help="Local root for saved adapter checkpoints")
    p.add_argument("--no-qlora", action="store_true",
                   help="Disable 4-bit QLoRA — needs 24 GB+ VRAM")
    p.add_argument("--dry-run", action="store_true",
                   help="Parse args and show dataset stats without training or pushing")
    args = p.parse_args()

    modes: list[str] = list(_GENERATORS.keys()) if args.train_all else \
                       ([args.mode] if args.mode else [])
    if not modes:
        p.error("Specify --mode <name> or --all")
    if args.data and len(modes) > 1:
        p.error("--data can only be combined with a single --mode")

    output_root = Path(args.output_dir)

    for mode in modes:
        train_one(
            mode=mode,
            base_model=args.base_model,
            hub_org=args.hub_org,
            n_examples=args.n,
            epochs=args.epochs,
            jsonl_path=args.data if len(modes) == 1 else None,
            output_root=output_root,
            use_qlora=not args.no_qlora,
            dry_run=args.dry_run,
        )

    if not args.dry_run:
        print("\nAll adapters pushed. Add these to your .env:")
        for mode in modes:
            print(f"  ADAPTER_{mode.upper()}={args.hub_org}/adapter-{mode}")


if __name__ == "__main__":
    main()