scripts/targate_repro.py

#!/usr/bin/env python3
"""
TarGATE reproduction — Target-Aware Data Selection via Token-Attenuation Gates.

Faithful reimplementation of the core method from ICML 2026 paper #3063
(OpenReview: xaqSrbGpPN):

1. Token-level Information Retention Ratio (IRR) gates attenuate FFN outputs.
2. Joint optimization: encourage high IRR on reference (target) data and
   low IRR on candidate/noise data (quality_classification_loss).
3. Score candidates by instance-level mean IRR; select top-k%.
4. Compare vs Random / Loss-based selection (noisy recovery + optional SFT).
5. Measure selection wall-clock and cross-model score transfer.

Official anonymous code: https://anonymous.4open.science/r/TarGATE-4008
"""

from __future__ import annotations

import argparse
import json
import os
import random
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


@dataclass
class ReproConfig:
    selector_model: str = "Qwen/Qwen2.5-0.5B-Instruct"
    transfer_model: str = "Qwen/Qwen2.5-1.5B-Instruct"
    seed: int = 42
    max_length: int = 256
    # Noisy scenario sizes (scaled vs paper: 8k faker + ~7k gsm8k train)
    n_gsm8k_candidate: int = 800
    n_gsm8k_reference: int = 80
    n_noise: int = 800
    n_eval: int = 200
    selection_pct: float = 0.10  # top-10% (paper often uses 5%)
    warmup_epochs: int = 1
    warmup_lr: float = 3e-4
    quality_loss_weight: float = 1.0
    quality_loss_balance: float = 0.2
    batch_size: int = 4
    grad_accum: int = 4
    sft_epochs: int = 1
    sft_lr: float = 2e-5
    sft_max_steps: int = 80
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    output_dir: str = "outputs"
    run_sft: bool = True
    run_transfer: bool = True
    dtype: str = "bf16"


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# ---------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------


def make_noise_samples(n: int, seed: int = 42) -> List[Dict]:
    """Synthetic noise approximating Faker-style irrelevant instruction data."""
    rng = random.Random(seed)
    topics = [
        "recipe for purple soup",
        "inventory of imaginary insects",
        "weather on Neptune last Tuesday",
        "biography of a fictional cat mayor",
        "instructions to fold invisible origami",
        "list of unused color names",
        "diary of a sentient toaster",
        "rules for underwater chess",
    ]
    out = []
    for i in range(n):
        topic = topics[i % len(topics)]
        filler = " ".join(rng.choice(["lorem", "ipsum", "dolor", "sit", "amet", "noise"]) for _ in range(rng.randint(20, 60)))
        out.append(
            {
                "dataset": "noise",
                "id": f"noise_{i:06d}",
                "is_target": 0,
                "messages": [
                    {"role": "user", "content": f"Please write about {topic}. {filler}"},
                    {"role": "assistant", "content": f"Sure. Here is content about {topic}. {filler}"},
                ],
            }
        )
    return out


def load_gsm8k(split: str = "train") -> List[Dict]:
    from datasets import load_dataset

    ds = load_dataset("openai/gsm8k", "main", split=split)
    rows = []
    for i, ex in enumerate(ds):
        rows.append(
            {
                "dataset": "gsm8k",
                "id": f"gsm8k_{split}_{i:06d}",
                "is_target": 1,
                "messages": [
                    {"role": "user", "content": ex["question"]},
                    {"role": "assistant", "content": ex["answer"]},
                ],
            }
        )
    return rows


def split_gsm8k(
    n_cand: int, n_ref: int, n_eval: int, seed: int
) -> Tuple[List[Dict], List[Dict], List[Dict]]:
    rows = load_gsm8k("train")
    rng = random.Random(seed)
    rng.shuffle(rows)
    ref = rows[:n_ref]
    cand = rows[n_ref : n_ref + n_cand]
    # holdout for quick exact-match eval after SFT (from remaining train as proxy)
    rest = rows[n_ref + n_cand :]
    eval_rows = rest[:n_eval]
    for r in ref:
        r["is_target"] = 1
    for r in cand:
        r["is_target"] = 1
    return cand, ref, eval_rows


def messages_to_text(messages: List[Dict]) -> str:
    return "\n".join(f"{m['role']}: {m['content']}" for m in messages)


class ChatDataset(Dataset):
    def __init__(self, rows: List[Dict], tokenizer, max_length: int, for_lm: bool = True):
        self.rows = rows
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.for_lm = for_lm

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx: int) -> Dict:
        row = self.rows[idx]
        text = messages_to_text(row["messages"])
        enc = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding=False,
            return_tensors=None,
        )
        item = {
            "input_ids": enc["input_ids"],
            "attention_mask": enc["attention_mask"],
            "is_target_task": int(row.get("is_target", 0)),
            "dataset": row["dataset"],
            "id": row["id"],
            "messages": row["messages"],
        }
        if self.for_lm:
            item["labels"] = list(enc["input_ids"])
        return item


def collate_fn(batch: List[Dict], pad_id: int) -> Dict:
    max_len = max(len(b["input_ids"]) for b in batch)

    def pad(seq, value):
        return seq + [value] * (max_len - len(seq))

    input_ids = torch.tensor([pad(b["input_ids"], pad_id) for b in batch], dtype=torch.long)
    attention_mask = torch.tensor([pad(b["attention_mask"], 0) for b in batch], dtype=torch.long)
    out = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "is_target_task": torch.tensor([b["is_target_task"] for b in batch], dtype=torch.long),
        "dataset": [b["dataset"] for b in batch],
        "id": [b["id"] for b in batch],
        "messages": [b["messages"] for b in batch],
    }
    if "labels" in batch[0]:
        labels = torch.tensor([pad(b["labels"], -100) for b in batch], dtype=torch.long)
        out["labels"] = labels
    return out


# ---------------------------------------------------------------------------
# TarGATE: inject token-attenuation gates into transformer MLP residual path
# ---------------------------------------------------------------------------


class QualityGate(nn.Module):
    def __init__(self, hidden_size: int, init_mean: float = 0.0, init_std: float = 0.02):
        super().__init__()
        self.gate = nn.Linear(hidden_size, 1, bias=False)
        nn.init.normal_(self.gate.weight, mean=init_mean, std=init_std)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # hidden: [B, T, H] -> good_ratio [B, T, 1]
        return torch.sigmoid(self.gate(hidden_states))


class GatedMLP(nn.Module):
    """MLP wrapper: out = original_mlp(x) * sigmoid(gate(x)). Must be nn.Module for HF layers."""

    def __init__(self, original_mlp: nn.Module, gate: QualityGate):
        super().__init__()
        self.original_mlp = original_mlp
        self.gate = gate
        self.capture = False
        self.last_ratio: Optional[torch.Tensor] = None

    def forward(self, hidden_states, *args, **kwargs):
        # Ensure gate weights match activation dtype (bf16/fp16)
        if next(self.gate.parameters()).dtype != hidden_states.dtype:
            self.gate.to(dtype=hidden_states.dtype)
        good_ratio = self.gate(hidden_states)
        mlp_out = self.original_mlp(hidden_states, *args, **kwargs)
        if self.capture:
            self.last_ratio = good_ratio.detach()
        return mlp_out * good_ratio.to(dtype=mlp_out.dtype)


class TarGATEWrapper(nn.Module):
    """
    Wrap a HF CausalLM and attenuate each decoder layer's MLP residual update
    by a learned token-level IRR gate (as in QualityGateDecoderLayer).
    """

    def __init__(self, model: nn.Module, quality_loss_weight: float = 1.0, quality_loss_balance: float = 0.2):
        super().__init__()
        self.model = model
        self.quality_loss_weight = quality_loss_weight
        self.quality_loss_balance = quality_loss_balance
        self.gates = nn.ModuleList()
        self.gated_mlps = nn.ModuleList()
        self._handles = []
        self._last_ratios: List[torch.Tensor] = []
        self._capture = False

        layers = self._get_layers()
        hidden = model.config.hidden_size
        # Match base model parameter dtype/device (e.g. bf16 on GPU)
        ref_param = next(model.parameters())
        for layer in layers:
            g = QualityGate(hidden)
            g.to(device=ref_param.device, dtype=ref_param.dtype)
            self.gates.append(g)
            # Patch MLP residual: residual + mlp(x) * good_ratio
            self._patch_layer(layer, g)

        # freeze base, train only gates
        for p in self.model.parameters():
            p.requires_grad = False
        for g in self.gates:
            for p in g.parameters():
                p.requires_grad = True

    def _get_layers(self):
        if hasattr(self.model, "model") and hasattr(self.model.model, "layers"):
            return self.model.model.layers
        if hasattr(self.model, "transformer") and hasattr(self.model.transformer, "h"):
            return self.model.transformer.h
        raise RuntimeError("Unsupported model architecture for TarGATE wrapping")

    def _patch_layer(self, layer: nn.Module, gate: QualityGate) -> None:
        if not hasattr(layer, "mlp"):
            raise RuntimeError("Layer has no mlp module")
        original_mlp = layer.mlp
        gated = GatedMLP(original_mlp, gate)
        self.gated_mlps.append(gated)
        # object.__setattr__ avoids Module type checks if needed; GatedMLP is a Module so OK
        layer.mlp = gated
        layer._targate_gate = gate

    def train(self, mode: bool = True):
        super().train(mode)
        # base stays frozen for BN-like behavior; keep eval for dropout stability
        self.model.eval()
        for g in self.gates:
            g.train(mode)
        return self

    def eval(self):
        super().eval()
        self.model.eval()
        for g in self.gates:
            g.eval()
        return self

    def quality_loss(
        self,
        ratios: List[torch.Tensor],
        attention_mask: torch.Tensor,
        is_target_task: torch.Tensor,
    ) -> torch.Tensor:
        """Match quality_classification_loss from official QualityGate code.

        For candidate samples (is_target=0): push mean IRR down.
        For reference samples (is_target=1): push mean IRR up.
        """
        if not ratios:
            return torch.tensor(0.0, device=attention_mask.device, requires_grad=True)
        total = None
        n = 0
        mask = attention_mask.float()
        is_train_t = 1.0 - is_target_task.float()
        for good_ratio in ratios:
            gr = good_ratio.squeeze(-1)  # [B, T]
            avg_per = (gr * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
            loss_train_t = is_train_t * self.quality_loss_balance * avg_per
            loss_target_t = (1.0 - is_train_t) * (2.0 - self.quality_loss_balance) * (1.0 - avg_per)
            layer_loss = (loss_train_t + loss_target_t).mean()
            total = layer_loss if total is None else total + layer_loss
            n += 1
        return total / max(n, 1)

    def forward(self, input_ids, attention_mask=None, labels=None, is_target_task=None, capture_gates=False, **kwargs):
        self._last_ratios = []
        self._capture = bool(capture_gates or (self.training and is_target_task is not None))
        for gm in self.gated_mlps:
            gm.capture = self._capture
            gm.last_ratio = None
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs)
        loss = outputs.loss
        qloss = None
        if is_target_task is not None:
            # Differentiable IRR path via post-attn LN hooks
            live_ratios = self._compute_live_ratios(input_ids, attention_mask)
            qloss = self.quality_loss_weight * self.quality_loss(live_ratios, attention_mask, is_target_task)
            if loss is not None:
                loss = loss + qloss
            else:
                loss = qloss
        return outputs, loss, qloss

    def _compute_live_ratios(self, input_ids, attention_mask) -> List[torch.Tensor]:
        """Forward hidden states layer-by-layer to get live gate outputs (differentiable)."""
        model = self.model
        embeds = model.model.embed_tokens(input_ids)
        hidden = embeds
        # Build position ids
        position_ids = attention_mask.long().cumsum(-1) - 1
        position_ids.masked_fill_(attention_mask == 0, 1)
        # Use model rotary if available via layer forward — simpler: call full model with hooks
        ratios = []
        # Use hooks on post_attention_layernorm outputs
        captured = []

        def make_hook(gate):
            def hook(module, inp, out):
                captured.append((gate, out))

            return hook

        handles = []
        layers = self._get_layers()
        for layer, gate in zip(layers, self.gates):
            if hasattr(layer, "post_attention_layernorm"):
                handles.append(layer.post_attention_layernorm.register_forward_hook(make_hook(gate)))
            elif hasattr(layer, "ln_2"):
                handles.append(layer.ln_2.register_forward_hook(make_hook(gate)))

        with torch.set_grad_enabled(self.training):
            _ = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)

        for h in handles:
            h.remove()

        for gate, h_states in captured:
            ratios.append(gate(h_states))
        return ratios

    def score_batch(self, input_ids, attention_mask) -> torch.Tensor:
        """Instance-level mean IRR across tokens and layers."""
        self.eval()
        with torch.no_grad():
            ratios = self._compute_live_ratios(input_ids, attention_mask)
            if not ratios:
                return torch.zeros(input_ids.size(0), device=input_ids.device)
            stacked = torch.stack([r.squeeze(-1) for r in ratios], dim=0)  # L,B,T
            mask = attention_mask.unsqueeze(0).float()
            summed = (stacked * mask).sum(dim=(0, 2))
            denom = mask.sum(dim=(0, 2)).clamp(min=1) * stacked.size(0) / stacked.size(0)
            # denom: sum over layers and tokens
            denom = attention_mask.float().sum(dim=1).clamp(min=1) * stacked.size(0)
            scores = summed / denom
            return scores


def load_base_model(model_name: str, device: str, dtype: str):
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch_dtype = torch.bfloat16 if dtype == "bf16" and torch.cuda.is_available() else torch.float32
    tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch_dtype,
        trust_remote_code=True,
    )
    model.to(device)
    return model, tok


# ---------------------------------------------------------------------------
# Baselines
# ---------------------------------------------------------------------------


@torch.no_grad()
def score_by_nll(model, tokenizer, rows: List[Dict], cfg: ReproConfig) -> List[float]:
    """Lower NLL = higher quality score (we invert)."""
    model.eval()
    ds = ChatDataset(rows, tokenizer, cfg.max_length, for_lm=True)
    loader = DataLoader(
        ds,
        batch_size=cfg.batch_size,
        shuffle=False,
        collate_fn=lambda b: collate_fn(b, tokenizer.pad_token_id),
    )
    scores = []
    for batch in loader:
        input_ids = batch["input_ids"].to(cfg.device)
        attention_mask = batch["attention_mask"].to(cfg.device)
        labels = batch["labels"].to(cfg.device)
        out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        # per-example CE
        logits = out.logits
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        loss_fct = nn.CrossEntropyLoss(reduction="none", ignore_index=-100)
        token_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        token_loss = token_loss.view(shift_labels.size())
        mask = (shift_labels != -100).float()
        per = (token_loss * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        # quality score = -NLL so higher is better (easier / more fluent)
        scores.extend((-per).detach().cpu().tolist())
    return scores


def score_random(rows: List[Dict], seed: int) -> List[float]:
    rng = random.Random(seed)
    return [rng.random() for _ in rows]


@torch.no_grad()
def score_targate(wrapper: TarGATEWrapper, tokenizer, rows: List[Dict], cfg: ReproConfig) -> List[float]:
    wrapper.eval()
    ds = ChatDataset(rows, tokenizer, cfg.max_length, for_lm=False)
    loader = DataLoader(
        ds,
        batch_size=cfg.batch_size,
        shuffle=False,
        collate_fn=lambda b: collate_fn(b, tokenizer.pad_token_id),
    )
    scores = []
    for batch in loader:
        input_ids = batch["input_ids"].to(cfg.device)
        attention_mask = batch["attention_mask"].to(cfg.device)
        s = wrapper.score_batch(input_ids, attention_mask)
        scores.extend(s.detach().cpu().tolist())
    return scores


def select_top_pct(rows: List[Dict], scores: List[float], pct: float) -> List[Dict]:
    k = max(1, int(np.ceil(pct * len(rows))))
    order = np.argsort(-np.array(scores))
    selected = []
    for idx in order[:k]:
        r = dict(rows[int(idx)])
        r["score"] = float(scores[int(idx)])
        selected.append(r)
    return selected


def selection_metrics(selected: List[Dict], candidate: List[Dict]) -> Dict:
    n_sel = len(selected)
    n_target = sum(1 for r in selected if r.get("dataset") == "gsm8k" or r.get("is_target") == 1)
    n_noise = sum(1 for r in selected if r.get("dataset") == "noise")
    pure = n_target / max(n_sel, 1)
    # recovery: fraction of all gsm8k candidates recovered
    n_cand_target = sum(1 for r in candidate if r.get("dataset") == "gsm8k")
    recovery = n_target / max(n_cand_target, 1)
    return {
        "n_selected": n_sel,
        "n_target_in_selected": n_target,
        "n_noise_in_selected": n_noise,
        "precision": pure,
        "recovery": recovery,
        "candidate_size": len(candidate),
        "candidate_target": n_cand_target,
    }


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------


def train_gates(
    wrapper: TarGATEWrapper,
    tokenizer,
    candidate: List[Dict],
    reference: List[Dict],
    cfg: ReproConfig,
) -> Dict:
    """Stage 1 warmup: joint LM + quality loss on ref+cand."""
    # subsample candidate for warmup (paper subset_ratio=0.25)
    rng = random.Random(cfg.seed)
    warmup_cand = list(candidate)
    rng.shuffle(warmup_cand)
    warmup_cand = warmup_cand[: max(len(reference) * 4, int(0.5 * len(candidate)))]
    train_rows = reference + warmup_cand
    rng.shuffle(train_rows)

    ds = ChatDataset(train_rows, tokenizer, cfg.max_length, for_lm=True)
    loader = DataLoader(
        ds,
        batch_size=cfg.batch_size,
        shuffle=True,
        collate_fn=lambda b: collate_fn(b, tokenizer.pad_token_id),
    )
    opt = torch.optim.AdamW([p for p in wrapper.gates.parameters() if p.requires_grad], lr=cfg.warmup_lr)
    wrapper.train()
    t0 = time.time()
    step = 0
    losses = []
    for epoch in range(cfg.warmup_epochs):
        opt.zero_grad(set_to_none=True)
        for batch in tqdm(loader, desc=f"warmup_epoch{epoch}"):
            input_ids = batch["input_ids"].to(cfg.device)
            attention_mask = batch["attention_mask"].to(cfg.device)
            labels = batch["labels"].to(cfg.device)
            is_target = batch["is_target_task"].to(cfg.device)
            _, loss, qloss = wrapper(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=labels,
                is_target_task=is_target,
                capture_gates=True,
            )
            (loss / cfg.grad_accum).backward()
            if (step + 1) % cfg.grad_accum == 0:
                opt.step()
                opt.zero_grad(set_to_none=True)
            losses.append(float(loss.detach().cpu()))
            step += 1
    if step % cfg.grad_accum != 0:
        opt.step()
    wall = time.time() - t0
    return {
        "warmup_steps": step,
        "warmup_wall_sec": wall,
        "warmup_loss_mean": float(np.mean(losses[-50:])) if losses else None,
        "n_trainable_gate_params": sum(p.numel() for p in wrapper.gates.parameters()),
    }


def quick_sft_and_eval(
    model_name: str,
    train_rows: List[Dict],
    eval_rows: List[Dict],
    cfg: ReproConfig,
    tag: str,
) -> Dict:
    """Lightweight LoRA-free full SFT on a small model, short steps; eval exact numeric match rate."""
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch_dtype = torch.bfloat16 if cfg.dtype == "bf16" and torch.cuda.is_available() else torch.float32
    tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch_dtype, trust_remote_code=True)
    model.to(cfg.device)
    # train all params lightly — for 0.5B short steps is ok on GPU
    for p in model.parameters():
        p.requires_grad = True
    opt = torch.optim.AdamW(model.parameters(), lr=cfg.sft_lr)
    ds = ChatDataset(train_rows, tok, cfg.max_length, for_lm=True)
    loader = DataLoader(
        ds,
        batch_size=max(1, cfg.batch_size // 2),
        shuffle=True,
        collate_fn=lambda b: collate_fn(b, tok.pad_token_id),
    )
    model.train()
    t0 = time.time()
    step = 0
    losses = []
    while step < cfg.sft_max_steps:
        for batch in loader:
            input_ids = batch["input_ids"].to(cfg.device)
            attention_mask = batch["attention_mask"].to(cfg.device)
            labels = batch["labels"].to(cfg.device)
            out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            out.loss.backward()
            opt.step()
            opt.zero_grad(set_to_none=True)
            losses.append(float(out.loss.detach().cpu()))
            step += 1
            if step >= cfg.sft_max_steps:
                break
    sft_wall = time.time() - t0

    # eval: greedy generate short answers and check final number
    model.eval()
    correct = 0
    total = 0
    for row in eval_rows[: cfg.n_eval]:
        q = row["messages"][0]["content"]
        gold = row["messages"][1]["content"]
        gold_num = gold.split("####")[-1].strip() if "####" in gold else gold.strip().split()[-1]
        prompt = f"user: {q}\nassistant:"
        enc = tok(prompt, return_tensors="pt", truncation=True, max_length=cfg.max_length).to(cfg.device)
        with torch.no_grad():
            gen = model.generate(
                **enc,
                max_new_tokens=64,
                do_sample=False,
                pad_token_id=tok.pad_token_id,
            )
        text = tok.decode(gen[0][enc["input_ids"].shape[1] :], skip_special_tokens=True)
        pred_num = ""
        try:
            if "####" in text:
                tail = text.split("####")[-1].strip()
                parts = tail.split()
                pred_num = parts[0] if parts else ""
            else:
                parts = text.strip().split()
                pred_num = parts[-1] if parts else ""
        except Exception:
            pred_num = ""
        # normalize
        gold_num_n = "".join(c for c in str(gold_num) if c.isdigit() or c in ".-")
        pred_num_n = "".join(c for c in str(pred_num) if c.isdigit() or c in ".-")
        if gold_num_n and gold_num_n == pred_num_n:
            correct += 1
        total += 1
    acc = correct / max(total, 1)
    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return {
        "tag": tag,
        "sft_steps": step,
        "sft_wall_sec": sft_wall,
        "sft_loss_mean": float(np.mean(losses[-20:])) if losses else None,
        "eval_n": total,
        "eval_correct": correct,
        "eval_acc": acc,
        "n_train": len(train_rows),
    }


# ---------------------------------------------------------------------------
# Main experiment
# ---------------------------------------------------------------------------


def run(cfg: ReproConfig) -> Dict:
    set_seed(cfg.seed)
    out_dir = Path(cfg.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print("=== Loading GSM8K + noise ===")
    cand_gsm, ref, eval_rows = split_gsm8k(cfg.n_gsm8k_candidate, cfg.n_gsm8k_reference, cfg.n_eval, cfg.seed)
    noise = make_noise_samples(cfg.n_noise, cfg.seed)
    candidate = cand_gsm + noise
    random.Random(cfg.seed).shuffle(candidate)
    print(f"candidate={len(candidate)} (gsm8k={len(cand_gsm)}, noise={len(noise)}), ref={len(ref)}, eval={len(eval_rows)}")

    print(f"=== Loading selector model {cfg.selector_model} ===")
    base, tok = load_base_model(cfg.selector_model, cfg.device, cfg.dtype)
    wrapper = TarGATEWrapper(base, cfg.quality_loss_weight, cfg.quality_loss_balance)
    wrapper.to(cfg.device)

    print("=== Stage 1: Gate warmup ===")
    warmup_stats = train_gates(wrapper, tok, candidate, ref, cfg)
    print(warmup_stats)

    print("=== Stage 2: Scoring ===")
    t0 = time.time()
    scores_tg = score_targate(wrapper, tok, candidate, cfg)
    t_tg = time.time() - t0

    t0 = time.time()
    scores_nll = score_by_nll(wrapper.model, tok, candidate, cfg)
    t_nll = time.time() - t0

    t0 = time.time()
    scores_rnd = score_random(candidate, cfg.seed + 1)
    t_rnd = time.time() - t0

    methods = {
        "TarGATE": scores_tg,
        "NegNLL": scores_nll,
        "Random": scores_rnd,
    }
    timings = {
        "TarGATE_score_sec": t_tg,
        "NegNLL_score_sec": t_nll,
        "Random_score_sec": t_rnd,
        "TarGATE_warmup_sec": warmup_stats["warmup_wall_sec"],
    }

    results = {"selection": {}, "timings": timings, "warmup": warmup_stats, "config": asdict(cfg)}
    selected_sets = {}
    for name, sc in methods.items():
        sel = select_top_pct(candidate, sc, cfg.selection_pct)
        selected_sets[name] = sel
        metrics = selection_metrics(sel, candidate)
        results["selection"][name] = metrics
        print(f"{name}: precision={metrics['precision']:.3f} recovery={metrics['recovery']:.3f} n={metrics['n_selected']}")

    results["efficiency"] = {
        "trainable_gate_params": warmup_stats["n_trainable_gate_params"],
        "base_model": cfg.selector_model,
        "score_time_targate_sec": t_tg,
        "score_time_nll_sec": t_nll,
        "warmup_time_sec": warmup_stats["warmup_wall_sec"],
        "notes": "TarGATE trains only tiny linear gates; scoring is one forward pass per batch.",
    }
    # Persist after selection so SFT failures do not lose primary metrics
    with open(out_dir / "results_selection.json", "w") as f:
        json.dump(results, f, indent=2)
    with open(out_dir / "selected_targate.jsonl", "w") as f:
        for r in selected_sets["TarGATE"]:
            f.write(json.dumps(r) + "\n")
    with open(out_dir / "selected_random.jsonl", "w") as f:
        for r in selected_sets["Random"]:
            f.write(json.dumps(r) + "\n")
    print(f"Checkpointed selection results to {out_dir}")

    if cfg.run_transfer:
        print("=== Cross-model transfer (data selected by small model) ===")
        results["transfer"] = {
            "selector": cfg.selector_model,
            "target_model_for_sft": cfg.transfer_model if cfg.run_sft else None,
            "selected_by_small_model_precision": results["selection"]["TarGATE"]["precision"],
            "claim": "Small selector TarGATE curates high-purity target data usable for larger model SFT",
        }

    # Optional short SFT comparison on selector model
    if cfg.run_sft:
        print("=== Stage 3: short SFT comparison ===")
        sft_results = {}
        for name in ["TarGATE", "Random", "NegNLL"]:
            print(f"SFT on {name} selection...")
            try:
                sft_results[name] = quick_sft_and_eval(
                    cfg.selector_model, selected_sets[name], eval_rows, cfg, tag=name
                )
                print(sft_results[name])
            except Exception as e:
                print(f"SFT {name} failed: {e}")
                sft_results[name] = {"tag": name, "error": str(e)}
        results["sft"] = sft_results

        if cfg.run_transfer:
            print("=== Cross-model SFT on transfer model ===")
            xfer = {}
            for name in ["TarGATE", "Random"]:
                print(f"Transfer SFT {cfg.transfer_model} on {name}...")
                try:
                    xfer[name] = quick_sft_and_eval(
                        cfg.transfer_model, selected_sets[name], eval_rows, cfg, tag=f"xfer_{name}"
                    )
                    print(xfer[name])
                except Exception as e:
                    print(f"Transfer SFT {name} failed: {e}")
                    xfer[name] = {"tag": f"xfer_{name}", "error": str(e)}
            results["transfer_sft"] = xfer

    out_path = out_dir / "results.json"
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Wrote {out_path}")
    return results


def parse_args() -> ReproConfig:
    p = argparse.ArgumentParser()
    p.add_argument("--selector-model", default=ReproConfig.selector_model)
    p.add_argument("--transfer-model", default=ReproConfig.transfer_model)
    p.add_argument("--output-dir", default="outputs")
    p.add_argument("--n-gsm8k-candidate", type=int, default=800)
    p.add_argument("--n-gsm8k-reference", type=int, default=80)
    p.add_argument("--n-noise", type=int, default=800)
    p.add_argument("--n-eval", type=int, default=100)
    p.add_argument("--selection-pct", type=float, default=0.10)
    p.add_argument("--warmup-epochs", type=int, default=1)
    p.add_argument("--batch-size", type=int, default=4)
    p.add_argument("--sft-max-steps", type=int, default=80)
    p.add_argument("--max-length", type=int, default=256)
    p.add_argument("--no-sft", action="store_true")
    p.add_argument("--no-transfer", action="store_true")
    p.add_argument("--seed", type=int, default=42)
    args = p.parse_args()
    return ReproConfig(
        selector_model=args.selector_model,
        transfer_model=args.transfer_model,
        output_dir=args.output_dir,
        n_gsm8k_candidate=args.n_gsm8k_candidate,
        n_gsm8k_reference=args.n_gsm8k_reference,
        n_noise=args.n_noise,
        n_eval=args.n_eval,
        selection_pct=args.selection_pct,
        warmup_epochs=args.warmup_epochs,
        batch_size=args.batch_size,
        sft_max_steps=args.sft_max_steps,
        max_length=args.max_length,
        run_sft=not args.no_sft,
        run_transfer=not args.no_transfer,
        seed=args.seed,
    )


if __name__ == "__main__":
    cfg = parse_args()
    print(json.dumps(asdict(cfg), indent=2))
    results = run(cfg)
    print(json.dumps(results, indent=2, default=str))