"""SFT warmstart for the Cortex brain-selector router. Workstream B Phase 6b. Trains a LoRA adapter on the Phase 6a B3 corpus so the router starts GRPO already able to emit valid JSON: {"brain": "epi" | "logistics" | "governance"} This intentionally uses a tiny raw Transformers + PEFT loop instead of TRL's SFTTrainer. The objective is simple masked causal LM over the assistant completion tokens, and avoiding TRL keeps the HF Jobs training surface aligned with the no-TRL GRPO path in ``minimal_proof.py``. Local preflight: DRY_RUN=1 HF_TOKEN=dummy uv run python training/scripts/sft_warmstart_router.py Live run: HF_TOKEN=hf_xxx OUTPUT_REPO=Angshuman28/cortex-router-sft-warmstart \\ uv run python training/scripts/sft_warmstart_router.py """ from __future__ import annotations import json import os import random import re import sys import textwrap import time from typing import Any, Dict, List, Optional def _env(name: str, default: Optional[str] = None, *, required: bool = False) -> str: value = os.environ.get(name, default) if required and not value: raise SystemExit(f"[FATAL] env var {name} is required but unset") return value or "" DRY_RUN = _env("DRY_RUN", "0") not in ("0", "", "false", "False") HF_TOKEN = _env("HF_TOKEN", required=not DRY_RUN) MODEL_NAME = _env("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct") B3_DATASET_REPO = _env("B3_DATASET_REPO", "Angshuman28/crisisworld-b3-corpus") OUTPUT_REPO = _env("OUTPUT_REPO", "Angshuman28/cortex-router-sft-warmstart") OUTPUT_DIR = _env("OUTPUT_DIR", "/tmp/cortex_router_sft_lora") MAX_TRAIN_STEPS = int(_env("MAX_TRAIN_STEPS", "200")) PER_DEVICE_BATCH = int(_env("PER_DEVICE_BATCH", "4")) GRAD_ACCUM = int(_env("GRAD_ACCUM", "2")) LR = float(_env("LR", "2e-5")) LORA_RANK = int(_env("LORA_RANK", "16")) LORA_ALPHA = int(_env("LORA_ALPHA", str(LORA_RANK * 2))) LORA_DROPOUT = float(_env("LORA_DROPOUT", "0.05")) MAX_SEQ_LEN = int(_env("MAX_SEQ_LEN", "1536")) MAX_NEW_TOKENS = int(_env("MAX_NEW_TOKENS", "32")) SEED = int(_env("SEED", "42")) VALIDATION_ATTEMPTS = int(_env("VALIDATION_ATTEMPTS", "10")) VALIDATION_MIN_RATE = float(_env("VALIDATION_MIN_RATE", "0.8")) PUSH_TO_HUB = _env("PUSH_TO_HUB", "1") not in ("0", "", "false", "False") def log(*args: object) -> None: print("[router-sft]", *args, flush=True) ROUTER_SYSTEM_PROMPT = textwrap.dedent( """ You are the Cortex brain selector. Read one CrisisWorld observation and choose exactly one specialist brain to act next. Brain options: - epi: epidemiology, surveillance, case growth, and outbreak spread - logistics: scarce resources, hospital beds, mobile units, allocation - governance: legal constraints, compliance, escalation, restrictions Respond with exactly one JSON object and no prose: {"brain": "epi" | "logistics" | "governance"} """ ).strip() _BRAIN_ALIASES = { "epi": "epi", "epidemiology": "epi", "logistics": "logistics", "governance": "governance", } def parse_router_choice(raw_text: str) -> Optional[str]: text = raw_text.strip() text = re.sub(r"```(?:json)?\s*", "", text) text = re.sub(r"```\s*$", "", text).strip() candidates = [text] start = text.find("{") if start >= 0: depth = 0 for index, char in enumerate(text[start:], start): if char == "{": depth += 1 elif char == "}": depth -= 1 if depth == 0: candidates.append(text[start : index + 1]) break for candidate_text in candidates: try: candidate = json.loads(candidate_text) except json.JSONDecodeError: continue if isinstance(candidate, dict): brain = _BRAIN_ALIASES.get(str(candidate.get("brain", "")).strip().lower()) if brain is not None: return brain return None def router_completion(brain: str) -> str: canonical = _BRAIN_ALIASES.get(brain) if canonical is None: raise ValueError(f"unknown brain label {brain!r}") return json.dumps({"brain": canonical}, separators=(",", ":")) def render_prompt(tokenizer: Any, observation_text: str) -> str: return tokenizer.apply_chat_template( [ {"role": "system", "content": ROUTER_SYSTEM_PROMPT}, {"role": "user", "content": observation_text}, ], tokenize=False, add_generation_prompt=True, ) def preflight_dataset_access(dataset_repo: str, token: str) -> None: log(f"preflight: checking dataset {dataset_repo}") if DRY_RUN: log("DRY_RUN=1 - skipping Hub dataset access check") return from huggingface_hub import HfApi from huggingface_hub.utils import RepositoryNotFoundError try: HfApi().dataset_info(dataset_repo, token=token or None) except RepositoryNotFoundError as exc: raise SystemExit(f"[FATAL] dataset {dataset_repo} not found: {exc}") from exc log("preflight: dataset accessible") def make_collate(tokenizer: Any): def collate(features: List[Dict[str, List[int]]]) -> Dict[str, Any]: import torch max_len = max(len(item["input_ids"]) for item in features) input_ids, attention_mask, labels = [], [], [] for item in features: pad = max_len - len(item["input_ids"]) input_ids.append(item["input_ids"] + [tokenizer.pad_token_id] * pad) attention_mask.append([1] * len(item["input_ids"]) + [0] * pad) labels.append(item["labels"] + [-100] * pad) return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), } return collate def tokenize_rows(tokenizer: Any, rows: List[Dict[str, Any]]) -> List[Dict[str, List[int]]]: tokenized: List[Dict[str, List[int]]] = [] eos = tokenizer.eos_token or "" for row in rows: prompt = render_prompt(tokenizer, row["observation_text"]) completion = router_completion(row["deterministic_brain_choice"]) + eos prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] full_ids = tokenizer( prompt + completion, add_special_tokens=False, truncation=True, max_length=MAX_SEQ_LEN, )["input_ids"] labels = list(full_ids) prompt_cutoff = min(len(prompt_ids), len(labels)) labels[:prompt_cutoff] = [-100] * prompt_cutoff if all(label == -100 for label in labels): continue tokenized.append({"input_ids": full_ids, "labels": labels}) if not tokenized: raise SystemExit("[FATAL] no tokenized training rows; check corpus schema") return tokenized def validate_router_json( model: Any, tokenizer: Any, rows: List[Dict[str, Any]], device: Any ) -> float: import torch sample_rows = rows[: max(1, VALIDATION_ATTEMPTS)] ok = 0 model.eval() for row in sample_rows: prompt = render_prompt(tokenizer, row["observation_text"]) inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN).to( device ) prompt_len = inputs["input_ids"].shape[1] with torch.no_grad(): out = model.generate( **inputs, do_sample=True, temperature=0.2, max_new_tokens=MAX_NEW_TOKENS, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) text = tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True) ok += parse_router_choice(text) is not None return ok / len(sample_rows) def main() -> int: log(f"MODEL_NAME={MODEL_NAME}") log(f"B3_DATASET_REPO={B3_DATASET_REPO}") log(f"OUTPUT_REPO={OUTPUT_REPO}") log(f"MAX_TRAIN_STEPS={MAX_TRAIN_STEPS} LR={LR} LORA_RANK={LORA_RANK}") preflight_dataset_access(B3_DATASET_REPO, HF_TOKEN) if DRY_RUN: log("DRY_RUN=1 - preflight only; not loading dataset/model or training") return 0 import torch from accelerate import Accelerator from datasets import load_dataset from huggingface_hub import HfApi from peft import LoraConfig, TaskType, get_peft_model from torch.utils.data import DataLoader from transformers import AutoModelForCausalLM, AutoTokenizer random.seed(SEED) torch.manual_seed(SEED) accelerator = Accelerator(gradient_accumulation_steps=GRAD_ACCUM) log("loading dataset") dsdict = load_dataset(B3_DATASET_REPO, token=HF_TOKEN) train_split = dsdict["train"] if "train" in dsdict else dsdict[list(dsdict.keys())[0]] rows = [dict(row) for row in train_split] random.shuffle(rows) required = {"observation_text", "deterministic_brain_choice"} missing = required - set(train_split.column_names) if missing: raise SystemExit(f"[FATAL] dataset missing columns {sorted(missing)}") log("loading tokenizer/model") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN) if tokenizer.pad_token_id is None: tokenizer.pad_token = tokenizer.eos_token dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=HF_TOKEN, torch_dtype=dtype) model = get_peft_model( model, LoraConfig( r=LORA_RANK, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", task_type=TaskType.CAUSAL_LM, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], ), ) optimizer = torch.optim.AdamW(model.parameters(), lr=LR) tokenized = tokenize_rows(tokenizer, rows) loader = DataLoader( tokenized, batch_size=PER_DEVICE_BATCH, shuffle=True, collate_fn=make_collate(tokenizer), ) model, optimizer, loader = accelerator.prepare(model, optimizer, loader) log(f"training rows={len(tokenized)}") step = 0 while step < MAX_TRAIN_STEPS: for batch in loader: with accelerator.accumulate(model): outputs = model(**batch) loss = outputs.loss accelerator.backward(loss) optimizer.step() optimizer.zero_grad() if accelerator.sync_gradients: step += 1 if step % max(MAX_TRAIN_STEPS // 20, 1) == 0: log(f"step={step}/{MAX_TRAIN_STEPS} loss={float(loss.detach().cpu()):.4f}") if step >= MAX_TRAIN_STEPS: break accelerator.wait_for_everyone() if accelerator.is_main_process: unwrapped = accelerator.unwrap_model(model) rate = validate_router_json(unwrapped, tokenizer, rows, accelerator.device) log(f"router JSON validation rate={rate:.0%}") if rate < VALIDATION_MIN_RATE: raise SystemExit( f"[FATAL] router JSON validation {rate:.0%} < {VALIDATION_MIN_RATE:.0%}" ) log(f"saving LoRA adapter to {OUTPUT_DIR}") unwrapped.save_pretrained(OUTPUT_DIR) tokenizer.save_pretrained(OUTPUT_DIR) if PUSH_TO_HUB: log(f"pushing to https://huggingface.co/{OUTPUT_REPO}") api = HfApi() api.create_repo( OUTPUT_REPO, exist_ok=True, repo_type="model", private=False, token=HF_TOKEN ) api.upload_folder( folder_path=OUTPUT_DIR, repo_id=OUTPUT_REPO, repo_type="model", token=HF_TOKEN, ) log("done") return 0 if __name__ == "__main__": t0 = time.time() rc = main() log(f"elapsed={time.time() - t0:.1f}s") sys.exit(rc)