v5d: save best-epoch checkpoint during SFT

a674764 verified 7 months ago

19.6 kB

#!/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
    submit_only: bool = False
    min_investigation_steps: int = 0

    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


_ALL_INVESTIGATION_ACTIONS = [
    {"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"},
]


def _expert_actions(
    env: InvoiceGuardEnvironment,
    task_id: TaskID,
    max_investigation_steps: int = 9,
) -> 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
    investigation = _ALL_INVESTIGATION_ACTIONS[:max_investigation_steps]
    used_names = [a["action_type"] for a in investigation]
    evidence = list(dict.fromkeys([*used_names, *gt.acceptable_evidence]))
    return [
        *investigation,
        {
            "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,
        },
    ]


TRACE_LENGTHS = [3, 5, 7, 9]
SUBMIT_LOSS_WEIGHT = 5.0


def build_sft_examples(
    tokenizer,
    env: InvoiceGuardEnvironment,
    tasks: list[TaskID],
    max_prompt_tokens: int,
) -> list[dict]:
    examples: list[dict] = []
    for task_id in tasks:
        for n_inv in TRACE_LENGTHS:
            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, max_investigation_steps=n_inv):
                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,
                    "trace_inv_steps": n_inv,
                })
                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,
    weight: float = 1.0,
) -> 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() * weight


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:
        if cfg.use_4bit:
            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)
    if cfg.submit_only:
        examples = [ex for ex in examples if ex["action_type"] == "submit_final_resolution"]
        print(f"[setup] submit-only mode: filtered to {len(examples)} submit examples", flush=True)
    if cfg.min_investigation_steps > 0:
        before = len(examples)
        examples = [ex for ex in examples if ex.get("trace_inv_steps", 0) >= cfg.min_investigation_steps]
        print(f"[setup] min_investigation_steps={cfg.min_investigation_steps}: {before} -> {len(examples)} examples", flush=True)
    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
    best_eval_score = -1.0
    best_epoch_dir: Optional[Path] = None
    for epoch in range(cfg.num_epochs):
        random.shuffle(examples)
        total_loss = 0.0
        model.train()
        for i, ex in enumerate(examples, 1):
            w = SUBMIT_LOSS_WEIGHT if ex["action_type"] == "submit_final_resolution" else 1.0
            loss = completion_loss(model, ex["prompt_ids"], ex["completion_ids"], device, weight=w)
            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:
            eval_result = evaluate(epoch + 1)
            log(eval_result)
            score = eval_result.get("eval/avg_grader_score", 0.0)
            if score > best_eval_score:
                best_eval_score = score
                best_epoch_dir = artifact_dir / f"best_epoch_{epoch+1}"
                best_epoch_dir.mkdir(parents=True, exist_ok=True)
                model.save_pretrained(str(best_epoch_dir))
                tokenizer.save_pretrained(str(best_epoch_dir))
                print(f"[checkpoint] new best epoch {epoch+1}: score={score:.4f}", flush=True)

    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)

        if best_epoch_dir and best_epoch_dir.exists():
            best_repo = f"{repo_id}-best"
            create_repo(repo_id=best_repo, repo_type="model", exist_ok=True, private=False, token=token)
            HfApi(token=token).upload_folder(
                folder_path=str(best_epoch_dir),
                repo_id=best_repo,
                repo_type="model",
                token=token,
                commit_message=f"Best epoch checkpoint (score={best_eval_score:.4f})",
            )
            print(f"[push] best checkpoint -> https://huggingface.co/{best_repo}", 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")
    p.add_argument("--no-4bit", action="store_true")
    p.add_argument("--submit-only", action="store_true")
    p.add_argument("--min-investigation-steps", type=int, default=None)
    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
    if args.no_4bit:
        cfg.use_4bit = False
    if args.submit_only:
        cfg.submit_only = True
    if args.min_investigation_steps is not None:
        cfg.min_investigation_steps = args.min_investigation_steps
    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
	submit_only: bool = False
	min_investigation_steps: int = 0

	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


	_ALL_INVESTIGATION_ACTIONS = [
	{"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"},
	]


	def _expert_actions(
	env: InvoiceGuardEnvironment,
	task_id: TaskID,
	max_investigation_steps: int = 9,
	) -> 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
	investigation = _ALL_INVESTIGATION_ACTIONS[:max_investigation_steps]
	used_names = [a["action_type"] for a in investigation]
	evidence = list(dict.fromkeys([used_names, gt.acceptable_evidence]))
	return [
	*investigation,
	{
	"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,
	},
	]


	TRACE_LENGTHS = [3, 5, 7, 9]
	SUBMIT_LOSS_WEIGHT = 5.0


	def build_sft_examples(
	tokenizer,
	env: InvoiceGuardEnvironment,
	tasks: list[TaskID],
	max_prompt_tokens: int,
	) -> list[dict]:
	examples: list[dict] = []
	for task_id in tasks:
	for n_inv in TRACE_LENGTHS:
	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, max_investigation_steps=n_inv):
	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,
	"trace_inv_steps": n_inv,
	})
	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,
	weight: float = 1.0,
	) -> 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() * weight


	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:
	if cfg.use_4bit:
	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)
	if cfg.submit_only:
	examples = [ex for ex in examples if ex["action_type"] == "submit_final_resolution"]
	print(f"[setup] submit-only mode: filtered to {len(examples)} submit examples", flush=True)
	if cfg.min_investigation_steps > 0:
	before = len(examples)
	examples = [ex for ex in examples if ex.get("trace_inv_steps", 0) >= cfg.min_investigation_steps]
	print(f"[setup] min_investigation_steps={cfg.min_investigation_steps}: {before} -> {len(examples)} examples", flush=True)
	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
	best_eval_score = -1.0
	best_epoch_dir: Optional[Path] = None
	for epoch in range(cfg.num_epochs):
	random.shuffle(examples)
	total_loss = 0.0
	model.train()
	for i, ex in enumerate(examples, 1):
	w = SUBMIT_LOSS_WEIGHT if ex["action_type"] == "submit_final_resolution" else 1.0
	loss = completion_loss(model, ex["prompt_ids"], ex["completion_ids"], device, weight=w)
	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:
	eval_result = evaluate(epoch + 1)
	log(eval_result)
	score = eval_result.get("eval/avg_grader_score", 0.0)
	if score > best_eval_score:
	best_eval_score = score
	best_epoch_dir = artifact_dir / f"best_epoch_{epoch+1}"
	best_epoch_dir.mkdir(parents=True, exist_ok=True)
	model.save_pretrained(str(best_epoch_dir))
	tokenizer.save_pretrained(str(best_epoch_dir))
	print(f"[checkpoint] new best epoch {epoch+1}: score={score:.4f}", flush=True)

	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)

	if best_epoch_dir and best_epoch_dir.exists():
	best_repo = f"{repo_id}-best"
	create_repo(repo_id=best_repo, repo_type="model", exist_ok=True, private=False, token=token)
	HfApi(token=token).upload_folder(
	folder_path=str(best_epoch_dir),
	repo_id=best_repo,
	repo_type="model",
	token=token,
	commit_message=f"Best epoch checkpoint (score={best_eval_score:.4f})",
	)
	print(f"[push] best checkpoint -> https://huggingface.co/{best_repo}", 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")
	p.add_argument("--no-4bit", action="store_true")
	p.add_argument("--submit-only", action="store_true")
	p.add_argument("--min-investigation-steps", type=int, default=None)
	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
	if args.no_4bit:
	cfg.use_4bit = False
	if args.submit_only:
	cfg.submit_only = True
	if args.min_investigation_steps is not None:
	cfg.min_investigation_steps = args.min_investigation_steps
	return cfg


	if __name__ == "__main__":
	main()