Add InvoiceGuard SFT and merge tooling

training/eval_adapter.py

@@ -1,184 +1,184 @@
-#!/usr/bin/env python3
-# /// script
-# requires-python = ">=3.10"
-# dependencies = [
-#   "torch>=2.2",
-#   "transformers>=4.46",
-#   "peft>=0.13",
-#   "accelerate>=1.0",
-#   "bitsandbytes>=0.43; platform_system != 'Darwin'",
-#   "huggingface_hub>=0.26",
-#   "openenv-core[core]>=0.2.1",
-#   "pydantic>=2.6",
-#   "pydantic-settings>=2.0",
-#   "fastapi>=0.115",
-#   "uvicorn>=0.30",
-#   "python-dotenv",
-#   "openai>=1.40",
-# ]
-# ///
-"""Evaluate a LoRA adapter on InvoiceGuard tasks and upload JSON artifacts."""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import sys
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Optional
-
-
-def _hf_token() -> Optional[str]:
-    return os.environ.get("HF_TOKEN") or os.environ.get("API_TOKEN_HF")
-
-
-def _bootstrap_invoice_guard_path() -> Path:
-    code_dir = os.environ.get("INVOICEGUARD_CODE_DIR")
-    if code_dir and Path(code_dir).is_dir():
-        sys.path.insert(0, code_dir)
-        return Path(code_dir)
-
-    repo = os.environ.get("INVOICEGUARD_CODE_REPO")
-    if repo:
-        from huggingface_hub import snapshot_download
-
-        local = snapshot_download(repo_id=repo, repo_type="model", token=_hf_token())
-        sys.path.insert(0, local)
-        return Path(local)
-
-    here = Path(__file__).resolve().parent.parent
-    sys.path.insert(0, str(here))
-    return here
-
-
-_CODE_ROOT = _bootstrap_invoice_guard_path()
-
-import torch
-from huggingface_hub import HfApi
-from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-
-from server.invoice_guard_environment import InvoiceGuardEnvironment  # type: ignore
-from tasks import HARD_TASK_LIST, TASK_LIST  # type: ignore
-from training.rollout import rollout_episode  # type: ignore
-
-
-def _task_slice(name: str):
-    if name == "canonical":
-        return list(TASK_LIST)
-    if name == "hard":
-        return list(HARD_TASK_LIST)
-    return list(TASK_LIST) + list(HARD_TASK_LIST)
-
-
-def main() -> None:
-    p = argparse.ArgumentParser()
-    p.add_argument("--base-model", default=os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507"))
-    p.add_argument("--adapter-repo", required=True)
-    p.add_argument("--slice", choices=["canonical", "hard", "all"], default="all")
-    p.add_argument("--max-tasks", type=int, default=None)
-    p.add_argument("--max-new-tokens", type=int, default=96)
-    p.add_argument("--max-prompt-tokens", type=int, default=2048)
-    p.add_argument("--artifact-dir", default="/tmp/invoiceguard-adapter-eval")
-    args = p.parse_args()
-
-    token = _hf_token()
-    if not token:
-        raise RuntimeError("HF_TOKEN/API_TOKEN_HF is required for adapter eval upload.")
-
-    print(f"[setup] code_root={_CODE_ROOT}", flush=True)
-    print(f"[setup] base_model={args.base_model}", flush=True)
-    print(f"[setup] adapter_repo={args.adapter_repo}", flush=True)
-    print(f"[setup] cuda available={torch.cuda.is_available()}", flush=True)
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    quant_cfg = None
-    if torch.cuda.is_available():
-        quant_cfg = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_compute_dtype=dtype,
-        )
-
-    tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=True, token=token)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    base = AutoModelForCausalLM.from_pretrained(
-        args.base_model,
-        torch_dtype=dtype,
-        device_map="auto" if torch.cuda.is_available() else None,
-        low_cpu_mem_usage=True,
-        quantization_config=quant_cfg,
-        token=token,
-    )
-    base.config.pad_token_id = tokenizer.pad_token_id
-    base.config.use_cache = False
-    model = PeftModel.from_pretrained(base, args.adapter_repo, token=token)
-    model.eval()
-
-    tasks = _task_slice(args.slice)
-    if args.max_tasks is not None:
-        tasks = tasks[: args.max_tasks]
-
-    env = InvoiceGuardEnvironment()
-    rows = []
-    for i, task_id in enumerate(tasks, 1):
-        print(f"[eval] {i}/{len(tasks)} {task_id.value}", flush=True)
-        traj = rollout_episode(
-            model,
-            tokenizer,
-            env,
-            task_id,
-            temperature=0.0001,
-            top_p=1.0,
-            max_new_tokens=args.max_new_tokens,
-            max_prompt_tokens=args.max_prompt_tokens,
-            device=device,
-        )
-        rows.append({
-            "task_id": task_id.value,
-            "grader_score": traj.grader_score,
-            "cumulative_reward": traj.cumulative_reward,
-            "success": traj.success,
-            "n_steps": traj.n_steps,
-            "terminal_decision": traj.terminal_decision,
-            "actions": [step.completion_text for step in traj.steps],
-            "step_rewards": [step.reward for step in traj.steps],
-        })
-
-    summary = {
-        "run_finished_at": datetime.now(timezone.utc).isoformat(),
-        "base_model": args.base_model,
-        "adapter_repo": args.adapter_repo,
-        "slice": args.slice,
-        "n_tasks": len(rows),
-        "avg_grader_score": sum(r["grader_score"] for r in rows) / max(len(rows), 1),
-        "avg_cumulative_reward": sum(r["cumulative_reward"] for r in rows) / max(len(rows), 1),
-        "success_rate": sum(1.0 if r["success"] else 0.0 for r in rows) / max(len(rows), 1),
-        "avg_steps": sum(r["n_steps"] for r in rows) / max(len(rows), 1),
-    }
-
-    out_dir = Path(args.artifact_dir)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    (out_dir / "adapter_eval_results.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
-    (out_dir / "adapter_eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
-    print(json.dumps(summary, indent=2), flush=True)
-
-    HfApi(token=token).upload_folder(
-        folder_path=str(out_dir),
-        repo_id=args.adapter_repo,
-        repo_type="model",
-        path_in_repo=f"eval_artifacts/{args.slice}",
-        token=token,
-        commit_message=f"Add InvoiceGuard adapter eval results ({args.slice})",
-    )
-    print(f"[push] eval artifacts uploaded to https://huggingface.co/{args.adapter_repo}", flush=True)
-
-
-if __name__ == "__main__":
-    main()
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch>=2.2",
+#   "transformers>=4.46",
+#   "peft>=0.13",
+#   "accelerate>=1.0",
+#   "bitsandbytes>=0.43; platform_system != 'Darwin'",
+#   "huggingface_hub>=0.26",
+#   "openenv-core[core]>=0.2.1",
+#   "pydantic>=2.6",
+#   "pydantic-settings>=2.0",
+#   "fastapi>=0.115",
+#   "uvicorn>=0.30",
+#   "python-dotenv",
+#   "openai>=1.40",
+# ]
+# ///
+"""Evaluate a LoRA adapter on InvoiceGuard tasks and upload JSON artifacts."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+
+def _hf_token() -> Optional[str]:
+    return os.environ.get("HF_TOKEN") or os.environ.get("API_TOKEN_HF")
+
+
+def _bootstrap_invoice_guard_path() -> Path:
+    code_dir = os.environ.get("INVOICEGUARD_CODE_DIR")
+    if code_dir and Path(code_dir).is_dir():
+        sys.path.insert(0, code_dir)
+        return Path(code_dir)
+
+    repo = os.environ.get("INVOICEGUARD_CODE_REPO")
+    if repo:
+        from huggingface_hub import snapshot_download
+
+        local = snapshot_download(repo_id=repo, repo_type="model", token=_hf_token())
+        sys.path.insert(0, local)
+        return Path(local)
+
+    here = Path(__file__).resolve().parent.parent
+    sys.path.insert(0, str(here))
+    return here
+
+
+_CODE_ROOT = _bootstrap_invoice_guard_path()
+
+import torch
+from huggingface_hub import HfApi
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+from server.invoice_guard_environment import InvoiceGuardEnvironment  # type: ignore
+from tasks import HARD_TASK_LIST, TASK_LIST  # type: ignore
+from training.rollout import rollout_episode  # type: ignore
+
+
+def _task_slice(name: str):
+    if name == "canonical":
+        return list(TASK_LIST)
+    if name == "hard":
+        return list(HARD_TASK_LIST)
+    return list(TASK_LIST) + list(HARD_TASK_LIST)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--base-model", default=os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507"))
+    p.add_argument("--adapter-repo", required=True)
+    p.add_argument("--slice", choices=["canonical", "hard", "all"], default="all")
+    p.add_argument("--max-tasks", type=int, default=None)
+    p.add_argument("--max-new-tokens", type=int, default=96)
+    p.add_argument("--max-prompt-tokens", type=int, default=2048)
+    p.add_argument("--artifact-dir", default="/tmp/invoiceguard-adapter-eval")
+    args = p.parse_args()
+
+    token = _hf_token()
+    if not token:
+        raise RuntimeError("HF_TOKEN/API_TOKEN_HF is required for adapter eval upload.")
+
+    print(f"[setup] code_root={_CODE_ROOT}", flush=True)
+    print(f"[setup] base_model={args.base_model}", flush=True)
+    print(f"[setup] adapter_repo={args.adapter_repo}", flush=True)
+    print(f"[setup] cuda available={torch.cuda.is_available()}", flush=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    quant_cfg = None
+    if torch.cuda.is_available():
+        quant_cfg = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=dtype,
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=True, token=token)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    base = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=dtype,
+        device_map="auto" if torch.cuda.is_available() else None,
+        low_cpu_mem_usage=True,
+        quantization_config=quant_cfg,
+        token=token,
+    )
+    base.config.pad_token_id = tokenizer.pad_token_id
+    base.config.use_cache = False
+    model = PeftModel.from_pretrained(base, args.adapter_repo, token=token)
+    model.eval()
+
+    tasks = _task_slice(args.slice)
+    if args.max_tasks is not None:
+        tasks = tasks[: args.max_tasks]
+
+    env = InvoiceGuardEnvironment()
+    rows = []
+    for i, task_id in enumerate(tasks, 1):
+        print(f"[eval] {i}/{len(tasks)} {task_id.value}", flush=True)
+        traj = rollout_episode(
+            model,
+            tokenizer,
+            env,
+            task_id,
+            temperature=0.0001,
+            top_p=1.0,
+            max_new_tokens=args.max_new_tokens,
+            max_prompt_tokens=args.max_prompt_tokens,
+            device=device,
+        )
+        rows.append({
+            "task_id": task_id.value,
+            "grader_score": traj.grader_score,
+            "cumulative_reward": traj.cumulative_reward,
+            "success": traj.success,
+            "n_steps": traj.n_steps,
+            "terminal_decision": traj.terminal_decision,
+            "actions": [step.completion_text for step in traj.steps],
+            "step_rewards": [step.reward for step in traj.steps],
+        })
+
+    summary = {
+        "run_finished_at": datetime.now(timezone.utc).isoformat(),
+        "base_model": args.base_model,
+        "adapter_repo": args.adapter_repo,
+        "slice": args.slice,
+        "n_tasks": len(rows),
+        "avg_grader_score": sum(r["grader_score"] for r in rows) / max(len(rows), 1),
+        "avg_cumulative_reward": sum(r["cumulative_reward"] for r in rows) / max(len(rows), 1),
+        "success_rate": sum(1.0 if r["success"] else 0.0 for r in rows) / max(len(rows), 1),
+        "avg_steps": sum(r["n_steps"] for r in rows) / max(len(rows), 1),
+    }
+
+    out_dir = Path(args.artifact_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    (out_dir / "adapter_eval_results.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
+    (out_dir / "adapter_eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(json.dumps(summary, indent=2), flush=True)
+
+    HfApi(token=token).upload_folder(
+        folder_path=str(out_dir),
+        repo_id=args.adapter_repo,
+        repo_type="model",
+        path_in_repo=f"eval_artifacts/{args.slice}",
+        token=token,
+        commit_message=f"Add InvoiceGuard adapter eval results ({args.slice})",
+    )
+    print(f"[push] eval artifacts uploaded to https://huggingface.co/{args.adapter_repo}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/merge_adapter.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch>=2.2",
+#   "transformers>=4.46",
+#   "peft>=0.13",
+#   "accelerate>=1.0",
+#   "huggingface_hub>=0.26",
+#   "safetensors>=0.4",
+# ]
+# ///
+"""Merge a PEFT LoRA adapter into its base model and push the merged model."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Optional
+
+import torch
+from huggingface_hub import create_repo
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def _hf_token() -> Optional[str]:
+    return os.environ.get("HF_TOKEN") or os.environ.get("API_TOKEN_HF")
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--base-model", default=os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507"))
+    p.add_argument("--adapter-repo", required=True)
+    p.add_argument("--merged-repo", required=True)
+    args = p.parse_args()
+
+    token = _hf_token()
+    if not token:
+        raise RuntimeError("HF_TOKEN/API_TOKEN_HF is required to push merged model.")
+
+    print(f"[setup] base_model={args.base_model}", flush=True)
+    print(f"[setup] adapter_repo={args.adapter_repo}", flush=True)
+    print(f"[setup] merged_repo={args.merged_repo}", flush=True)
+    print(f"[setup] cuda available={torch.cuda.is_available()}", flush=True)
+
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    base = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=dtype,
+        device_map="auto" if torch.cuda.is_available() else None,
+        low_cpu_mem_usage=True,
+        token=token,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=True, token=token)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    base.config.pad_token_id = tokenizer.pad_token_id
+
+    peft_model = PeftModel.from_pretrained(base, args.adapter_repo, token=token)
+    merged = peft_model.merge_and_unload()
+
+    create_repo(repo_id=args.merged_repo, repo_type="model", exist_ok=True, private=False, token=token)
+    print("[push] pushing merged model", flush=True)
+    merged.push_to_hub(
+        args.merged_repo,
+        private=False,
+        safe_serialization=True,
+        token=token,
+        commit_message="Save merged InvoiceGuard model",
+    )
+    tokenizer.push_to_hub(
+        args.merged_repo,
+        private=False,
+        token=token,
+        commit_message="Save merged InvoiceGuard tokenizer",
+    )
+    print(f"[push] done -> https://huggingface.co/{args.merged_repo}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/train_sft.py

@@ -0,0 +1,441 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch>=2.2",
+#   "transformers>=4.46",
+#   "peft>=0.13",
+#   "accelerate>=1.0",
+#   "bitsandbytes>=0.43; platform_system != 'Darwin'",
+#   "huggingface_hub>=0.26",
+#   "trackio>=0.1.4",
+#   "openenv-core[core]>=0.2.1",
+#   "pydantic>=2.6",
+#   "pydantic-settings>=2.0",
+#   "fastapi>=0.115",
+#   "uvicorn>=0.30",
+#   "python-dotenv",
+#   "openai>=1.40",
+# ]
+# ///
+"""InvoiceGuard supervised trace fine-tuning backup run.
+
+This is intentionally separate from GRPO. It trains a LoRA adapter on
+environment-generated expert traces so we have a deterministic supervised
+fallback artifact while online RL jobs are running.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import sys
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+
+def _hf_token() -> Optional[str]:
+    return os.environ.get("HF_TOKEN") or os.environ.get("API_TOKEN_HF")
+
+
+def _bootstrap_invoice_guard_path() -> Path:
+    code_dir = os.environ.get("INVOICEGUARD_CODE_DIR")
+    if code_dir and Path(code_dir).is_dir():
+        sys.path.insert(0, code_dir)
+        return Path(code_dir)
+
+    repo = os.environ.get("INVOICEGUARD_CODE_REPO")
+    if repo:
+        from huggingface_hub import snapshot_download
+
+        local = snapshot_download(repo_id=repo, repo_type="model", token=_hf_token())
+        sys.path.insert(0, local)
+        return Path(local)
+
+    here = Path(__file__).resolve().parent.parent
+    sys.path.insert(0, str(here))
+    return here
+
+
+_CODE_ROOT = _bootstrap_invoice_guard_path()
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import HfApi, create_repo
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+from inference import SYSTEM_PROMPT, build_action, build_observation_prompt  # type: ignore
+from models import TaskID  # type: ignore
+from server.invoice_guard_environment import InvoiceGuardEnvironment  # type: ignore
+from tasks import HARD_TASK_LIST, TASK_LIST  # type: ignore
+from training.rollout import rollout_episode  # type: ignore
+
+
+@dataclass
+class SftConfig:
+    base_model: str = os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507")
+    hub_username: Optional[str] = os.environ.get("HF_USERNAME")
+    hub_model_id: str = os.environ.get("HUB_MODEL_ID", "invoiceguard-qwen3-4b-sft")
+    trackio_project: str = os.environ.get("TRACKIO_PROJECT", "invoiceguard-round2")
+    trackio_run_name: str = os.environ.get("TRACKIO_RUN_NAME", "qwen3-4b-sft")
+    artifact_dir: str = os.environ.get("ARTIFACT_DIR", "/tmp/invoiceguard-sft-artifacts")
+
+    seed: int = 42
+    num_epochs: int = 4
+    max_train_tasks: Optional[int] = None
+    eval_holdout_canonical: int = 3
+    eval_holdout_hard: int = 3
+    eval_every_epoch: bool = True
+
+    lr: float = 5e-5
+    grad_clip: float = 1.0
+    max_prompt_tokens: int = 2048
+    max_new_tokens: int = 96
+    bf16: bool = torch.cuda.is_available()
+    use_4bit: bool = True
+    gradient_checkpointing: bool = True
+
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+    lora_target_modules: tuple = ("q_proj", "k_proj", "v_proj", "o_proj")
+
+    push_to_hub: bool = True
+
+
+def split_tasks(cfg: SftConfig) -> tuple[list[TaskID], list[TaskID]]:
+    rng = random.Random(cfg.seed)
+    canonical = list(TASK_LIST)
+    hard = list(HARD_TASK_LIST)
+    rng.shuffle(canonical)
+    rng.shuffle(hard)
+    eval_tasks = (
+        canonical[: cfg.eval_holdout_canonical]
+        + hard[: cfg.eval_holdout_hard]
+    )
+    train_tasks = (
+        canonical[cfg.eval_holdout_canonical:]
+        + hard[cfg.eval_holdout_hard:]
+    )
+    if cfg.max_train_tasks is not None:
+        train_tasks = train_tasks[: cfg.max_train_tasks]
+    return train_tasks, eval_tasks
+
+
+def _expert_actions(env: InvoiceGuardEnvironment, task_id: TaskID) -> list[dict]:
+    case = getattr(env, "_case", None)
+    if case is None:
+        env.reset(task_id=task_id.value)
+        case = getattr(env, "_case", None)
+    assert case is not None
+    gt = case.ground_truth
+    evidence = list(dict.fromkeys([
+        "inspect_purchase_order",
+        "inspect_goods_receipt_note",
+        "inspect_invoice_line_items",
+        "inspect_vendor_profile",
+        "compare_quantity",
+        "compare_price",
+        "compare_totals",
+        "check_for_duplicate_invoice",
+        "inspect_policy_rules",
+        *gt.acceptable_evidence,
+    ]))
+    return [
+        {"action_type": "inspect_purchase_order"},
+        {"action_type": "inspect_goods_receipt_note"},
+        {"action_type": "inspect_invoice_line_items"},
+        {"action_type": "inspect_vendor_profile"},
+        {"action_type": "compare_quantity"},
+        {"action_type": "compare_price"},
+        {"action_type": "compare_totals"},
+        {"action_type": "check_for_duplicate_invoice"},
+        {"action_type": "inspect_policy_rules"},
+        {
+            "action_type": "submit_final_resolution",
+            "final_decision": gt.correct_decision.value,
+            "exception_type": gt.correct_exception_type.value,
+            "evidence_references": evidence,
+            "explanation": "Key findings: " + "; ".join(gt.key_findings[:3]),
+            "confidence": 0.9,
+        },
+    ]
+
+
+def build_sft_examples(
+    tokenizer,
+    env: InvoiceGuardEnvironment,
+    tasks: list[TaskID],
+    max_prompt_tokens: int,
+) -> list[dict]:
+    examples: list[dict] = []
+    for task_id in tasks:
+        obs = env.reset(task_id=task_id.value)
+        messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
+        for action_dict in _expert_actions(env, task_id):
+            user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
+            messages.append({"role": "user", "content": user_msg})
+            prompt_text = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+            completion_text = json.dumps(action_dict, ensure_ascii=False)
+            prompt_ids = tokenizer(
+                prompt_text,
+                return_tensors="pt",
+                add_special_tokens=False,
+                truncation=True,
+                max_length=max_prompt_tokens,
+            ).input_ids[0]
+            completion_ids = tokenizer(
+                completion_text,
+                return_tensors="pt",
+                add_special_tokens=False,
+            ).input_ids[0]
+            examples.append({
+                "task_id": task_id.value,
+                "action_type": action_dict["action_type"],
+                "prompt_ids": prompt_ids,
+                "completion_ids": completion_ids,
+                "completion_text": completion_text,
+            })
+            messages.append({"role": "assistant", "content": completion_text})
+            obs = env.step(build_action(action_dict))
+            if obs.done:
+                break
+    return examples
+
+
+def completion_loss(model, prompt_ids: torch.Tensor, completion_ids: torch.Tensor, device: torch.device) -> torch.Tensor:
+    input_ids = torch.cat([prompt_ids, completion_ids], dim=0).unsqueeze(0).to(device)
+    attention_mask = torch.ones_like(input_ids)
+    out = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
+    logits = out.logits[0, :-1, :]
+    targets = input_ids[0, 1:]
+    logprobs = F.log_softmax(logits.float(), dim=-1)
+    token_lp = logprobs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+    comp_len = completion_ids.shape[0]
+    return -token_lp[-comp_len:].mean()
+
+
+def main() -> None:
+    cfg = _parse_args()
+    token = _hf_token()
+    if cfg.push_to_hub and (not token or not cfg.hub_username):
+        raise RuntimeError("HF_TOKEN and HF_USERNAME are required when pushing SFT output.")
+
+    print(f"[setup] code_root={_CODE_ROOT}", flush=True)
+    print(f"[setup] base_model={cfg.base_model}", flush=True)
+    print(f"[setup] cuda available={torch.cuda.is_available()}", flush=True)
+
+    random.seed(cfg.seed)
+    torch.manual_seed(cfg.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(cfg.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.bfloat16 if cfg.bf16 else torch.float32
+
+    artifact_dir = Path(cfg.artifact_dir)
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+    metrics_path = artifact_dir / "sft_metrics.jsonl"
+    summary_path = artifact_dir / "sft_summary.json"
+
+    tokenizer = AutoTokenizer.from_pretrained(cfg.base_model, use_fast=True, token=token)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    quant_cfg = None
+    if cfg.use_4bit and torch.cuda.is_available():
+        quant_cfg = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=dtype,
+        )
+    base = AutoModelForCausalLM.from_pretrained(
+        cfg.base_model,
+        torch_dtype=dtype,
+        device_map="auto" if torch.cuda.is_available() else None,
+        low_cpu_mem_usage=True,
+        quantization_config=quant_cfg,
+        token=token,
+    )
+    base.config.pad_token_id = tokenizer.pad_token_id
+    base.config.use_cache = False
+    if cfg.gradient_checkpointing:
+        base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
+        base.gradient_checkpointing_enable()
+
+    lora_cfg = LoraConfig(
+        r=cfg.lora_r,
+        lora_alpha=cfg.lora_alpha,
+        lora_dropout=cfg.lora_dropout,
+        target_modules=list(cfg.lora_target_modules),
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(base, lora_cfg)
+    model.print_trainable_parameters()
+
+    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=cfg.lr)
+    env = InvoiceGuardEnvironment()
+    train_tasks, eval_tasks = split_tasks(cfg)
+    examples = build_sft_examples(tokenizer, env, train_tasks, cfg.max_prompt_tokens)
+    print(f"[setup] train_tasks={len(train_tasks)} eval_tasks={len(eval_tasks)} examples={len(examples)}", flush=True)
+
+    tracker = None
+    try:
+        import trackio
+        tracker = trackio.init(
+            project=cfg.trackio_project,
+            name=cfg.trackio_run_name,
+            config={
+                "base_model": cfg.base_model,
+                "hub_model_id": cfg.hub_model_id,
+                "num_epochs": cfg.num_epochs,
+                "n_train_tasks": len(train_tasks),
+                "n_eval_tasks": len(eval_tasks),
+                "n_examples": len(examples),
+                "lr": cfg.lr,
+                "lora_r": cfg.lora_r,
+            },
+        )
+        print("[setup] trackio initialised", flush=True)
+    except Exception as e:
+        print(f"[setup] trackio disabled: {e}", flush=True)
+
+    def log(row: dict) -> None:
+        with metrics_path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps({"time": datetime.now(timezone.utc).isoformat(), **row}) + "\n")
+        print(" | ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in row.items()), flush=True)
+        if tracker is not None:
+            try:
+                trackio.log(row, step=int(row.get("step", 0)))
+            except Exception:
+                pass
+
+    def evaluate(epoch: int) -> dict:
+        model.eval()
+        scores, rewards, successes, steps = [], [], [], []
+        for task_id in eval_tasks:
+            traj = rollout_episode(
+                model,
+                tokenizer,
+                env,
+                task_id,
+                temperature=0.0001,
+                top_p=1.0,
+                max_new_tokens=cfg.max_new_tokens,
+                max_prompt_tokens=cfg.max_prompt_tokens,
+                device=device,
+            )
+            scores.append(traj.grader_score)
+            rewards.append(traj.cumulative_reward)
+            successes.append(1.0 if traj.success else 0.0)
+            steps.append(traj.n_steps)
+        model.train()
+        return {
+            "step": epoch,
+            "eval/avg_grader_score": sum(scores) / max(len(scores), 1),
+            "eval/avg_cum_reward": sum(rewards) / max(len(rewards), 1),
+            "eval/success_rate": sum(successes) / max(len(successes), 1),
+            "eval/avg_steps": sum(steps) / max(len(steps), 1),
+        }
+
+    t_start = time.time()
+    global_step = 0
+    for epoch in range(cfg.num_epochs):
+        random.shuffle(examples)
+        total_loss = 0.0
+        model.train()
+        for i, ex in enumerate(examples, 1):
+            loss = completion_loss(model, ex["prompt_ids"], ex["completion_ids"], device)
+            loss.backward()
+            total_loss += float(loss.detach().item())
+            torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], cfg.grad_clip)
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+            global_step += 1
+            if i % 25 == 0:
+                log({"step": global_step, "train/epoch": epoch + 1, "train/example": i, "train/loss": total_loss / i})
+        log({"step": global_step, "train/epoch": epoch + 1, "train/loss": total_loss / max(len(examples), 1)})
+        if cfg.eval_every_epoch:
+            log(evaluate(epoch + 1))
+
+    summary = {
+        "run_finished_at": datetime.now(timezone.utc).isoformat(),
+        "base_model": cfg.base_model,
+        "hub_model_id": cfg.hub_model_id,
+        "num_epochs": cfg.num_epochs,
+        "train_tasks": [t.value for t in train_tasks],
+        "eval_tasks": [t.value for t in eval_tasks],
+        "n_examples": len(examples),
+        "wall_clock_s": round(time.time() - t_start, 2),
+    }
+    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    if cfg.push_to_hub and cfg.hub_username:
+        assert token is not None
+        repo_id = f"{cfg.hub_username}/{cfg.hub_model_id}"
+        create_repo(repo_id=repo_id, repo_type="model", exist_ok=True, private=False, token=token)
+        print(f"[push] pushing SFT adapter to {repo_id}", flush=True)
+        model.push_to_hub(repo_id, private=False, token=token, commit_message="Save InvoiceGuard SFT adapter")
+        tokenizer.push_to_hub(repo_id, private=False, token=token, commit_message="Save InvoiceGuard SFT tokenizer")
+        HfApi(token=token).upload_folder(
+            folder_path=str(artifact_dir),
+            repo_id=repo_id,
+            repo_type="model",
+            path_in_repo="sft_artifacts",
+            token=token,
+            commit_message="Add InvoiceGuard SFT artifacts",
+        )
+        print(f"[push] done -> https://huggingface.co/{repo_id}", flush=True)
+
+
+def _parse_args() -> SftConfig:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-name", dest="base_model", default=None)
+    p.add_argument("--hub-model-id", default=None)
+    p.add_argument("--num-epochs", type=int, default=None)
+    p.add_argument("--max-train-tasks", type=int, default=None)
+    p.add_argument("--eval-holdout-canonical", type=int, default=None)
+    p.add_argument("--eval-holdout-hard", type=int, default=None)
+    p.add_argument("--lr", type=float, default=None)
+    p.add_argument("--max-new-tokens", type=int, default=None)
+    p.add_argument("--max-prompt-tokens", type=int, default=None)
+    p.add_argument("--no-push", action="store_true")
+    args = p.parse_args()
+
+    cfg = SftConfig()
+    if args.base_model:
+        cfg.base_model = args.base_model
+    if args.hub_model_id:
+        cfg.hub_model_id = args.hub_model_id
+    if args.num_epochs is not None:
+        cfg.num_epochs = args.num_epochs
+    if args.max_train_tasks is not None:
+        cfg.max_train_tasks = args.max_train_tasks
+    if args.eval_holdout_canonical is not None:
+        cfg.eval_holdout_canonical = args.eval_holdout_canonical
+    if args.eval_holdout_hard is not None:
+        cfg.eval_holdout_hard = args.eval_holdout_hard
+    if args.lr is not None:
+        cfg.lr = args.lr
+    if args.max_new_tokens is not None:
+        cfg.max_new_tokens = args.max_new_tokens
+    if args.max_prompt_tokens is not None:
+        cfg.max_prompt_tokens = args.max_prompt_tokens
+    if args.no_push:
+        cfg.push_to_hub = False
+    return cfg
+
+
+if __name__ == "__main__":
+    main()