Upload folder using huggingface_hub

.env.example

@@ -0,0 +1,32 @@
+# -- InvoiceGuard Environment Variables --------------------------
+#
+# Copy this file to .env and fill in your values:
+#   cp .env.example .env
+#
+# These are read by inference.py via python-dotenv.
+# Never commit .env to git (it's already in .gitignore).
+# ----------------------------------------------------------------
+
+# 1. API_BASE_URL -- The LLM API endpoint.
+#    OpenAI:         https://api.openai.com/v1
+#    HF Inference:   https://router.huggingface.co/v1
+#    Local/custom:   http://localhost:8080/v1
+API_BASE_URL=https://api.openai.com/v1
+
+# 2. MODEL_NAME -- The model identifier for inference.
+#    OpenAI:  gpt-4.1-mini, gpt-5.4-mini, etc.
+#    HF:     Qwen/Qwen2.5-72B-Instruct, meta-llama/Llama-3.1-8B-Instruct, etc.
+MODEL_NAME=gpt-4.1-mini
+
+# 3. HF_TOKEN -- Your Hugging Face / API key.
+#    This is the PRIMARY key used by the inference script for LLM calls.
+#    For OpenAI models: put your OpenAI key here (sk-...)
+#    For HF models: put your HF token here (hf_...)
+#    Get HF token at: https://huggingface.co/settings/tokens
+HF_TOKEN=
+
+# 5. LOCAL_IMAGE_NAME -- Docker image name for containerized mode.
+#    Only needed if running inference against a Docker container
+#    instead of the local environment directly.
+#    Build first: docker build -t invoiceguard .
+# LOCAL_IMAGE_NAME=invoiceguard

inference.py

@@ -168,6 +168,39 @@ def strip_think_blocks(text: str) -> str:
     return _THINK_RE.sub("", text).strip()
 
 
+def _extract_first_json_object(text: str) -> dict | None:
+    """Find the first balanced {...} in text and parse it."""
+    start = text.find("{")
+    if start == -1:
+        return None
+    depth = 0
+    in_str = False
+    escape = False
+    for i in range(start, len(text)):
+        c = text[i]
+        if escape:
+            escape = False
+            continue
+        if c == "\\":
+            escape = True
+            continue
+        if c == '"':
+            in_str = not in_str
+            continue
+        if in_str:
+            continue
+        if c == "{":
+            depth += 1
+        elif c == "}":
+            depth -= 1
+            if depth == 0:
+                try:
+                    return json.loads(text[start : i + 1])
+                except json.JSONDecodeError:
+                    return None
+    return None
+
+
 def parse_llm_response(response_text: str) -> dict:
     """Extract a JSON object from the LLM response."""
     text = strip_think_blocks(response_text).strip()
@@ -190,6 +223,10 @@ def parse_llm_response(response_text: str) -> dict:
             except json.JSONDecodeError:
                 continue
 
+    obj = _extract_first_json_object(text)
+    if obj is not None:
+        return obj
+
     return {"action_type": "summarize_findings"}
 
 

training/merge_adapter.py

@@ -1,82 +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()
+#!/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/rollout.py

@@ -133,6 +133,12 @@ def rollout_episode(
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
+        if traj.n_steps < 2:
+            print(f"[rollout-diag] task={task_id.value} step={traj.n_steps} "
+                  f"gen_tokens={len(completion_ids)} "
+                  f"raw_text={repr(raw_text[:300])} "
+                  f"completion_text={repr(completion_text[:200])}", flush=True)
+
         messages.append({"role": "assistant", "content": completion_text})
 
         params = parse_llm_response(completion_text)

training/train_grpo.py

@@ -355,11 +355,16 @@ def run_format_warmup(
                 truncation=True,
                 max_length=cfg.max_prompt_tokens,
             ).input_ids[0]
-            completion_ids = tokenizer(
+            comp_enc = tokenizer(
                 completion_text,
                 return_tensors="pt",
                 add_special_tokens=False,
             ).input_ids[0]
+            eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+            if eos_id is not None and eos_id != tokenizer.unk_token_id:
+                completion_ids = torch.cat([comp_enc, torch.tensor([eos_id])])
+            else:
+                completion_ids = comp_enc
             lp = _completion_logprobs(policy, prompt_ids, completion_ids, device)
             loss = -lp / max(int(completion_ids.shape[0]), 1)
             loss.backward()

training/train_sft.py

@@ -1,445 +1,450 @@
-#!/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 = 384
-    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})
-            try:
-                prompt_text = tokenizer.apply_chat_template(
-                    messages, tokenize=False, add_generation_prompt=True,
-                    enable_thinking=False,
-                )
-            except TypeError:
-                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()
+#!/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 = 384
+    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})
+            try:
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True,
+                    enable_thinking=False,
+                )
+            except TypeError:
+                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]
+            comp_enc = tokenizer(
+                completion_text,
+                return_tensors="pt",
+                add_special_tokens=False,
+            ).input_ids[0]
+            eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+            if eos_id is not None and eos_id != tokenizer.unk_token_id:
+                completion_ids = torch.cat([comp_enc, torch.tensor([eos_id])])
+            else:
+                completion_ids = comp_enc
+            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()