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