Add training support package

detection_tokens/__init__.py

@@ -0,0 +1,21 @@
+"""Core package for training and using detection-token checkpoints."""
+
+from .config import PipelineConfig
+from .pipeline import (
+    prepare_data,
+    run_dataset_binary_evaluation,
+    run_holdout_evaluation,
+    run_training,
+    run_training_pipeline,
+    score_single_text,
+)
+
+__all__ = [
+    "PipelineConfig",
+    "prepare_data",
+    "run_dataset_binary_evaluation",
+    "run_holdout_evaluation",
+    "run_training",
+    "run_training_pipeline",
+    "score_single_text",
+]

detection_tokens/checkpoints.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+
+@dataclass(slots=True)
+class TokenCheckpoint:
+    new_token: str
+    token_id: int | None
+    embedding: torch.Tensor
+    loss_history: list[float]
+    secondary_embeddings: list[torch.Tensor] | None = None
+
+
+def save_token_checkpoint(
+    *,
+    token: str,
+    token_id: int,
+    embedding: torch.Tensor,
+    loss_history: list[float],
+    path: Path,
+    secondary_embeddings: list[torch.Tensor] | None = None,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "new_token": token,
+            "token_id": token_id,
+            "embedding": embedding.detach().cpu(),
+            "loss_history": list(loss_history),
+            "secondary_embeddings": (
+                [row.detach().cpu() for row in secondary_embeddings]
+                if secondary_embeddings
+                else None
+            ),
+        },
+        path,
+    )
+
+
+def load_token_checkpoint(path: Path) -> TokenCheckpoint:
+    payload = torch.load(path, map_location="cpu", weights_only=True)
+    token = payload.get("new_token", payload.get("token"))
+    if token is None:
+        raise KeyError(f"Checkpoint {path} has no token metadata")
+    return TokenCheckpoint(
+        new_token=str(token),
+        token_id=payload.get("token_id"),
+        embedding=payload["embedding"].detach().cpu(),
+        loss_history=list(payload.get("loss_history", [])),
+        secondary_embeddings=[
+            row.detach().cpu()
+            for row in (payload.get("secondary_embeddings") or [])
+        ] or None,
+    )

detection_tokens/config.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+def model_slug(model_name: str) -> str:
+    return model_name.replace("/", "__")
+
+
+@dataclass(slots=True)
+class PangramBinaryConfig:
+    enabled: bool = True
+    dataset_name: str = "pangram/editlens_iclr"
+    dataset_split: str = "train"
+    local_dataset_path: Path | None = Path("/home/ubuntu/data/pangram_editlens_iclr")
+    ai_text_types: tuple[str, ...] = ("ai_generated",)
+    human_text_types: tuple[str, ...] = ("human_written",)
+    train_pairs: int = 5_000
+
+
+@dataclass(slots=True)
+class RaidBinaryConfig:
+    enabled: bool = True
+    dataset_name: str = "liamdugan/raid"
+    dataset_split: str = "train"
+    human_model_name: str = "human"
+    require_attack_none: bool = True
+    train_pairs: int = 5_000
+    eval_holdout_pairs: int = 1_000
+
+
+@dataclass(slots=True)
+class DataConfig:
+    task_name: str = "binary_human_vs_ai"
+    training_holdout_pairs: int = 1_000
+    min_text_chars: int = 200
+    pangram: PangramBinaryConfig = field(default_factory=PangramBinaryConfig)
+    raid: RaidBinaryConfig = field(default_factory=RaidBinaryConfig)
+
+
+@dataclass(slots=True)
+class ModelConfig:
+    model_name: str = "meta-llama/Llama-3.1-8B-Instruct"
+    ai_token: str = "<ai>"
+    human_token: str = "<human>"
+    max_length: int = 384
+    prompt_template: str = "Write {token} text."
+
+
+@dataclass(slots=True)
+class TrainingConfig:
+    ai_learning_rate: float = 1.0e-4
+    human_learning_rate: float = 1.0e-4
+    batch_size: int = 16
+    train_steps: int | None = 625
+    beta: float = 0.1
+    apo_alpha: float = 1.0
+    warmup_steps: int = 20
+    min_learning_rate: float = 1.0e-6
+    ref_cache_batch_size: int = 16
+    eval_subset_size: int = 128
+    eval_every_steps: int = 50
+    neutral_word: str | None = None  # if set, warm-start new token rows from this vocab token
+    source_balance_by_source: bool = False
+
+
+@dataclass(slots=True)
+class CheckpointInitConfig:
+    ai_token_path: Path | None = None
+    human_token_path: Path | None = None
+
+
+@dataclass(slots=True)
+class EvaluationConfig:
+    mode: str = "holdout"
+    saved_pairs_filename: str = "holdout_pairs.json"
+    binary_split: str = "test"
+    positive_text_types: tuple[str, ...] = ("ai_generated",)
+    negative_text_types: tuple[str, ...] = ("human_written",)
+    binary_output_path: Path | None = None
+
+
+@dataclass(slots=True)
+class ScoringConfig:
+    text: str = "Example text to score."
+    score_mode: str = "avg_margin"
+    token_sigmoid_tau: float = 1.0
+
+
+@dataclass(slots=True)
+class VerbalizationTokenSetConfig:
+    name: str
+    token_dir: Path
+
+
+def default_verbalization_token_sets() -> tuple[VerbalizationTokenSetConfig, ...]:
+    root = Path("/lambda/nfs/daniel-dc/detectiontokens_outputs")
+    return (
+        VerbalizationTokenSetConfig(
+            name="raid_only",
+            token_dir=root / "raid_only_train_raid_holdout_eval" / "tokens",
+        ),
+        VerbalizationTokenSetConfig(
+            name="pangram_only",
+            token_dir=root / "pangram_1k_train_pangram_test_eval" / "tokens",
+        ),
+    )
+
+
+@dataclass(slots=True)
+class VerbalizationConfig:
+    output_dir: Path = Path("/lambda/nfs/daniel-dc/detectiontokens_outputs/verbalizations")
+    token_sets: tuple[VerbalizationTokenSetConfig, ...] = field(default_factory=default_verbalization_token_sets)
+    n_samples: int = 5
+    max_new_tokens: int = 128
+
+
+@dataclass(slots=True)
+class OutputConfig:
+    output_root: Path = Path("/lambda/nfs/daniel-dc/detectiontokens_outputs/pangram_raid_binary_llama31_8b")
+    repo_root: Path = Path("/lambda/nfs/daniel-dc/DetectionTokens")
+
+    @property
+    def splits_dir(self) -> Path:
+        return self.output_root / "splits"
+
+    @property
+    def tokens_dir(self) -> Path:
+        return self.output_root / "tokens"
+
+    @property
+    def model_tokens_root(self) -> Path:
+        return self.repo_root / "tokens"
+
+    def model_tokens_dir(self, model_name: str) -> Path:
+        return self.model_tokens_root / model_slug(model_name)
+
+    @property
+    def cache_dir(self) -> Path:
+        return self.output_root / "ref_cache"
+
+    @property
+    def evaluation_dir(self) -> Path:
+        return self.output_root / "evaluation"
+
+
+@dataclass(slots=True)
+class PipelineConfig:
+    seed: int = 42
+    data: DataConfig = field(default_factory=DataConfig)
+    model: ModelConfig = field(default_factory=ModelConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    init_checkpoints: CheckpointInitConfig = field(default_factory=CheckpointInitConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    scoring: ScoringConfig = field(default_factory=ScoringConfig)
+    verbalization: VerbalizationConfig = field(default_factory=VerbalizationConfig)
+    output: OutputConfig = field(default_factory=OutputConfig)

detection_tokens/data.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import json
+import random
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+from datasets import Dataset, load_dataset
+
+from .config import DataConfig
+
+
+@dataclass(slots=True)
+class SourcePair:
+    pair_id: str
+    text_id: str
+    source_id: str
+    dataset_name: str
+    source: str
+    model: str
+    text_type: str
+    cosine_score: float | None
+    ai_text: str
+    human_text: str
+
+
+@dataclass(slots=True)
+class BinaryEvalRow:
+    row_id: str
+    text: str
+    label: int
+    text_type: str
+    model: str
+    source_id: str
+
+
+def _valid_text(text: str, *, min_text_chars: int) -> bool:
+    return isinstance(text, str) and len(text.strip()) >= min_text_chars
+
+
+def _normalize_text(text: Any) -> str:
+    return str(text or "").strip()
+
+
+def _find_local_arrow_file(root: Path, split: str) -> Path:
+    direct_path = root / f"editlens_iclr-{split}.arrow"
+    if direct_path.exists():
+        return direct_path
+    matches = sorted(root.rglob(f"editlens_iclr-{split}.arrow"))
+    if not matches:
+        raise FileNotFoundError(f"Missing local dataset file for split={split!r} under {root}")
+    return matches[0]
+
+
+def _pick_best_row(rows: list[dict[str, Any]], *, text_key: str) -> dict[str, Any]:
+    if not rows:
+        raise ValueError("Cannot pick from an empty row list.")
+    rows = sorted(
+        rows,
+        key=lambda row: (
+            _normalize_text(row.get("prompt")) == "",
+            _normalize_text(row.get("title")) == "",
+            _normalize_text(row.get(text_key)) == "",
+        ),
+    )
+    return rows[0]
+
+
+def load_pangram_rows(config: DataConfig, *, split: str) -> Dataset:
+    if config.pangram.local_dataset_path is not None:
+        arrow_path = _find_local_arrow_file(config.pangram.local_dataset_path, split)
+        return Dataset.from_file(str(arrow_path))
+    return load_dataset(config.pangram.dataset_name, split=split)
+
+
+def load_raid_rows(config: DataConfig, *, split: str) -> Dataset:
+    return load_dataset(config.raid.dataset_name, split=split)
+
+
+def build_pangram_binary_pairs(config: DataConfig) -> list[SourcePair]:
+    rows = load_pangram_rows(config, split=config.pangram.dataset_split)
+    ai_rows = []
+    human_by_text_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
+
+    for row in rows:
+        text = _normalize_text(row.get("text"))
+        if not _valid_text(text, min_text_chars=config.min_text_chars):
+            continue
+        text_type = _normalize_text(row.get("text_type"))
+        if text_type in config.pangram.human_text_types:
+            text_id = _normalize_text(row.get("text_id"))
+            source_id = _normalize_text(row.get("source_id"))
+            if text_id:
+                human_by_text_id[text_id].append(dict(row))
+            if source_id:
+                human_by_source_id[source_id].append(dict(row))
+        elif text_type in config.pangram.ai_text_types:
+            ai_rows.append(dict(row))
+
+    pairs: list[SourcePair] = []
+    for row in ai_rows:
+        pair_source_id = _normalize_text(row.get("source_id"))
+        if not pair_source_id:
+            continue
+        candidates = human_by_text_id.get(pair_source_id)
+        if not candidates:
+            candidates = human_by_source_id.get(pair_source_id)
+        if not candidates:
+            continue
+        human_row = _pick_best_row(candidates, text_key="text")
+        pairs.append(
+            SourcePair(
+                pair_id=f"pangram::{pair_source_id}::{_normalize_text(row.get('text_id'))}",
+                text_id=_normalize_text(row.get("text_id")),
+                source_id=pair_source_id,
+                dataset_name="pangram",
+                source=_normalize_text(row.get("source")),
+                model=_normalize_text(row.get("model")),
+                text_type=_normalize_text(row.get("text_type")),
+                cosine_score=float(row["cosine_score"]) if row.get("cosine_score") is not None else None,
+                ai_text=_normalize_text(row.get("text")),
+                human_text=_normalize_text(human_row.get("text")),
+            )
+        )
+    return pairs
+
+
+def build_raid_binary_pairs(config: DataConfig) -> list[SourcePair]:
+    rows = load_raid_rows(config, split=config.raid.dataset_split)
+    human_rows: list[dict[str, Any]] = []
+    ai_rows: list[dict[str, Any]] = []
+    for row in rows:
+        if config.raid.require_attack_none and _normalize_text(row.get("attack")) not in {"", "none"}:
+            continue
+        text = _normalize_text(row.get("generation"))
+        if not _valid_text(text, min_text_chars=config.min_text_chars):
+            continue
+        model = _normalize_text(row.get("model"))
+        if model == config.raid.human_model_name:
+            human_rows.append(dict(row))
+        else:
+            ai_rows.append(dict(row))
+
+    human_by_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    human_by_source_id: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for row in human_rows:
+        row_id = _normalize_text(row.get("id"))
+        source_id = _normalize_text(row.get("source_id"))
+        if row_id:
+            human_by_id[row_id].append(row)
+        if source_id:
+            human_by_source_id[source_id].append(row)
+
+    pairs: list[SourcePair] = []
+    for row in ai_rows:
+        pair_source_id = _normalize_text(row.get("source_id"))
+        if not pair_source_id:
+            continue
+        candidates = human_by_id.get(pair_source_id)
+        if not candidates:
+            candidates = human_by_source_id.get(pair_source_id)
+        if not candidates:
+            continue
+        human_row = _pick_best_row(candidates, text_key="generation")
+        pairs.append(
+            SourcePair(
+                pair_id=f"raid::{pair_source_id}::{_normalize_text(row.get('model'))}::{_normalize_text(row.get('id'))}",
+                text_id=_normalize_text(row.get("id")),
+                source_id=pair_source_id,
+                dataset_name="raid",
+                source=_normalize_text(row.get("domain")),
+                model=_normalize_text(row.get("model")),
+                text_type="ai_generated",
+                cosine_score=None,
+                ai_text=_normalize_text(row.get("generation")),
+                human_text=_normalize_text(human_row.get("generation")),
+            )
+        )
+    return pairs
+
+
+def _take_pairs(
+    pairs: list[SourcePair],
+    *,
+    take: int,
+    seed: int,
+) -> tuple[list[SourcePair], list[SourcePair]]:
+    rng = random.Random(seed)
+    shuffled = list(pairs)
+    rng.shuffle(shuffled)
+    if len(shuffled) < take:
+        raise ValueError(f"Need at least {take} pairs, found {len(shuffled)}.")
+    return shuffled[:take], shuffled[take:]
+
+
+def build_training_and_eval_splits(
+    config: DataConfig,
+    *,
+    seed: int,
+) -> tuple[list[SourcePair], list[SourcePair], list[SourcePair], dict[str, int]]:
+    source_pools: dict[str, list[SourcePair]] = {}
+    if config.pangram.enabled:
+        source_pools["pangram"] = build_pangram_binary_pairs(config)
+    if config.raid.enabled:
+        source_pools["raid"] = build_raid_binary_pairs(config)
+
+    train_pairs: list[SourcePair] = []
+    holdout_candidates: list[SourcePair] = []
+    raid_eval_pairs: list[SourcePair] = []
+    metadata = {f"{name}_pairs_available": len(pairs) for name, pairs in source_pools.items()}
+
+    if config.raid.enabled:
+        raid_eval_pairs, remaining_raid = _take_pairs(
+            source_pools["raid"],
+            take=config.raid.eval_holdout_pairs,
+            seed=seed + 100,
+        )
+        source_pools["raid"] = remaining_raid
+
+    if config.pangram.enabled:
+        selected, remaining = _take_pairs(
+            source_pools["pangram"],
+            take=config.pangram.train_pairs,
+            seed=seed + 1,
+        )
+        train_pairs.extend(selected)
+        holdout_candidates.extend(remaining)
+
+    if config.raid.enabled:
+        selected, remaining = _take_pairs(
+            source_pools["raid"],
+            take=config.raid.train_pairs,
+            seed=seed + 2,
+        )
+        train_pairs.extend(selected)
+        holdout_candidates.extend(remaining)
+
+    holdout_pairs, _ = _take_pairs(
+        holdout_candidates,
+        take=config.training_holdout_pairs,
+        seed=seed + 3,
+    )
+    random.Random(seed + 4).shuffle(train_pairs)
+    random.Random(seed + 5).shuffle(holdout_pairs)
+    metadata.update(
+        {
+            "train_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in train_pairs),
+            "train_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in train_pairs),
+            "holdout_pairs_from_pangram": sum(pair.dataset_name == "pangram" for pair in holdout_pairs),
+            "holdout_pairs_from_raid": sum(pair.dataset_name == "raid" for pair in holdout_pairs),
+            "raid_eval_pairs_from_raid": len(raid_eval_pairs),
+        }
+    )
+    return train_pairs, holdout_pairs, raid_eval_pairs, metadata
+
+
+def save_pairs(path: Path, pairs: list[SourcePair]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps([asdict(pair) for pair in pairs], indent=2), encoding="utf-8")
+
+
+def load_pairs(path: Path) -> list[SourcePair]:
+    rows = json.loads(path.read_text(encoding="utf-8"))
+    return [SourcePair(**row) for row in rows]
+
+
+def load_binary_eval_rows(
+    config: DataConfig,
+    *,
+    split: str,
+    positive_text_types: set[str],
+    negative_text_types: set[str],
+) -> list[BinaryEvalRow]:
+    rows = load_pangram_rows(config, split=split)
+    payload: list[BinaryEvalRow] = []
+    for index, row in enumerate(rows):
+        text_type = str(row.get("text_type", "")).strip()
+        if text_type in positive_text_types:
+            label = 1
+        elif text_type in negative_text_types:
+            label = 0
+        else:
+            continue
+        text = str(row.get("text", "")).strip()
+        if not _valid_text(text, min_text_chars=config.min_text_chars):
+            continue
+        payload.append(
+            BinaryEvalRow(
+                row_id=str(row.get("text_id", index)),
+                text=text,
+                label=label,
+                text_type=text_type,
+                model=str(row.get("model", "")),
+                source_id=str(row.get("source_id", "")),
+            )
+        )
+    return payload

detection_tokens/modeling.py

@@ -0,0 +1,263 @@
+from __future__ import annotations
+
+import math
+import os
+import importlib.util
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import login
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+from .checkpoints import load_token_checkpoint
+from .config import PipelineConfig
+
+
+def _load_causal_lm(model_name: str, **kwargs):
+    """Load a causal LM, falling back to architecture-specific classes when
+    AutoModelForCausalLM doesn't recognize the model type."""
+    try:
+        return AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+    except (ValueError, ModuleNotFoundError):
+        config = AutoConfig.from_pretrained(model_name)
+        model_type = getattr(config, "model_type", "")
+        if model_type == "gemma3":
+            from transformers import Gemma3ForConditionalGeneration
+            return Gemma3ForConditionalGeneration.from_pretrained(model_name, **kwargs)
+        if model_type == "gemma4":
+            from transformers import Gemma4ForConditionalGeneration
+            return Gemma4ForConditionalGeneration.from_pretrained(model_name, **kwargs)
+        raise
+
+
+def resolve_hf_token() -> str | None:
+    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+
+
+def maybe_login_hf() -> None:
+    token = resolve_hf_token()
+    if token:
+        login(token=token, add_to_git_credential=False)
+
+
+def resolve_dtype() -> torch.dtype:
+    return torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+
+def has_accelerate() -> bool:
+    return importlib.util.find_spec("accelerate") is not None
+
+
+@dataclass(slots=True)
+class ModelBundle:
+    config: PipelineConfig
+    tokenizer: AutoTokenizer
+    model: AutoModelForCausalLM
+    initial_tokenizer_len: int
+
+
+def initialize_model_bundle(config: PipelineConfig) -> ModelBundle:
+    maybe_login_hf()
+    hf_token = resolve_hf_token()
+    dtype = resolve_dtype()
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model.model_name, use_fast=True, token=hf_token)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    initial_len = len(tokenizer)
+    tokenizer.add_special_tokens(
+        {"additional_special_tokens": [config.model.ai_token, config.model.human_token]}
+    )
+
+    load_kwargs = {
+        "token": hf_token,
+        "dtype": dtype,
+    }
+    if torch.cuda.is_available() and has_accelerate():
+        load_kwargs["device_map"] = "auto"
+    model = _load_causal_lm(config.model.model_name, **load_kwargs)
+    if torch.cuda.is_available() and not has_accelerate():
+        model = model.to("cuda")
+    elif not torch.cuda.is_available():
+        model = model.to("cpu")
+
+    try:
+        model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
+    except TypeError:
+        model.resize_token_embeddings(len(tokenizer))
+
+    # Mean-fill Gemma 4's auxiliary per-layer token table. Some Transformers
+    # versions resize it internally, while others leave it at the old size.
+    for name, module in model.named_modules():
+        if isinstance(module, torch.nn.Embedding) and module is not model.get_input_embeddings():
+            if module.weight.shape[0] == initial_len:
+                mean_row = module.weight.data[:initial_len].mean(
+                    dim=0, dtype=torch.float32
+                ).to(dtype=module.weight.dtype)
+                new_emb = torch.nn.Embedding(
+                    len(tokenizer), module.weight.shape[1],
+                    device=module.weight.device, dtype=module.weight.dtype,
+                )
+                new_emb.weight.data[:initial_len] = module.weight.data
+                new_emb.weight.data[initial_len:] = mean_row
+                if not torch.equal(
+                    new_emb.weight.data[initial_len], mean_row
+                ):
+                    raise RuntimeError(
+                        f"Failed to mean-fill resized embedding {name}"
+                    )
+                parent_name, attr_name = name.rsplit(".", 1)
+                parent = dict(model.named_modules())[parent_name]
+                setattr(parent, attr_name, new_emb)
+                print(
+                    f"Mean-filled secondary embedding {name}: "
+                    f"{initial_len} -> {len(tokenizer)}"
+                )
+            elif module.weight.shape[0] == len(tokenizer):
+                mean_row = module.weight.data[:initial_len].mean(
+                    dim=0, dtype=torch.float32
+                ).to(dtype=module.weight.dtype)
+                module.weight.data[initial_len:] = mean_row
+                if not torch.equal(module.weight.data[initial_len], mean_row):
+                    raise RuntimeError(
+                        f"Failed to mean-fill expanded embedding {name}"
+                    )
+                print(
+                    f"Mean-filled expanded secondary embedding {name}: "
+                    f"rows {initial_len}:{len(tokenizer)}"
+                )
+
+    model.config.use_cache = False
+    if hasattr(model, "gradient_checkpointing_enable"):
+        model.gradient_checkpointing_enable()
+    model.eval()
+
+    input_emb = model.get_input_embeddings()
+    with torch.no_grad():
+        mean_in = input_emb.weight[:initial_len].mean(dim=0)
+        for token in (config.model.ai_token, config.model.human_token):
+            token_id = tokenizer.convert_tokens_to_ids(token)
+            input_emb.weight[token_id].copy_(mean_in + torch.randn_like(mean_in) * 1e-5)
+
+    bundle = ModelBundle(
+        config=config,
+        tokenizer=tokenizer,
+        model=model,
+        initial_tokenizer_len=initial_len,
+    )
+    apply_initial_checkpoints(bundle)
+    return bundle
+
+
+def apply_initial_checkpoints(bundle: ModelBundle) -> None:
+    input_emb = bundle.model.get_input_embeddings()
+    token_dir = bundle.config.output.model_tokens_dir(bundle.config.model.model_name)
+    ai_path = bundle.config.init_checkpoints.ai_token_path or token_dir / "ai_token.pt"
+    human_path = bundle.config.init_checkpoints.human_token_path or token_dir / "human_token.pt"
+
+    secondary_embs = [
+        module
+        for module in bundle.model.modules()
+        if (
+            isinstance(module, torch.nn.Embedding)
+            and module is not input_emb
+            and module.weight.shape[0] == len(bundle.tokenizer)
+        )
+    ]
+
+    def install(path, token):
+        if not path.exists():
+            return
+        checkpoint = load_token_checkpoint(path)
+        token_id = bundle.tokenizer.convert_tokens_to_ids(token)
+        saved_secondary = checkpoint.secondary_embeddings or []
+        if len(saved_secondary) != len(secondary_embs):
+            raise ValueError(
+                f"{path} has {len(saved_secondary)} secondary rows, but "
+                f"{bundle.config.model.model_name} exposes "
+                f"{len(secondary_embs)} secondary token embeddings."
+            )
+        input_emb.weight[token_id].copy_(
+            checkpoint.embedding.to(
+                input_emb.weight.device, dtype=input_emb.weight.dtype
+            )
+        )
+        for embedding, row in zip(secondary_embs, saved_secondary):
+            embedding.weight[token_id].copy_(
+                row.to(embedding.weight.device, dtype=embedding.weight.dtype)
+            )
+
+    with torch.no_grad():
+        install(ai_path, bundle.config.model.ai_token)
+        install(human_path, bundle.config.model.human_token)
+
+
+def build_prompt(bundle: ModelBundle, token: str) -> str:
+    content = bundle.config.model.prompt_template.format(token=token)
+    messages = [{"role": "user", "content": content}]
+    return bundle.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+
+def encode_response(bundle: ModelBundle, prompt_text: str, response_text: str) -> tuple[torch.Tensor, int]:
+    full_text = prompt_text + response_text
+    prompt_ids = bundle.tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+    full_ids = bundle.tokenizer(
+        full_text,
+        return_tensors="pt",
+        add_special_tokens=False,
+        truncation=True,
+        max_length=bundle.config.model.max_length,
+    )["input_ids"][0]
+    return full_ids, int(len(prompt_ids))
+
+
+def compute_average_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
+    token_logps = compute_token_logprobs(bundle, input_ids, prompt_len)
+    return token_logps.mean()
+
+
+def _model_forward(bundle: ModelBundle, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None):
+    """Run a forward pass, injecting token_type_ids for models that require it."""
+    fwd_kwargs: dict = {"input_ids": input_ids}
+    if attention_mask is not None:
+        fwd_kwargs["attention_mask"] = attention_mask
+    model_type = getattr(bundle.model.config, "model_type", "")
+    if model_type in ("gemma3", "gemma4"):
+        fwd_kwargs["token_type_ids"] = torch.zeros_like(input_ids)
+    return bundle.model(**fwd_kwargs)
+
+
+def compute_token_logprobs(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
+    input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
+    logits = _model_forward(bundle, input_ids).logits[0]
+    shift_logits = logits[prompt_len - 1 : -1]
+    shift_labels = input_ids[0, prompt_len:]
+    log_probs = F.log_softmax(shift_logits, dim=-1)
+    return log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]
+
+
+def compute_sequence_logprob(bundle: ModelBundle, input_ids: torch.Tensor, prompt_len: int) -> torch.Tensor:
+    input_ids = input_ids.unsqueeze(0).to(bundle.model.device)
+    logits = _model_forward(bundle, input_ids).logits[0]
+    shift_logits = logits[prompt_len - 1 : -1]
+    shift_labels = input_ids[0, prompt_len:]
+    log_probs = F.log_softmax(shift_logits, dim=-1)
+    token_logps = log_probs[torch.arange(len(shift_labels), device=bundle.model.device), shift_labels]
+    return token_logps.sum()
+
+
+def cosine_with_floor(
+    step: int,
+    total_steps: int,
+    base_lr: float,
+    *,
+    min_lr: float,
+    warmup_steps: int,
+) -> float:
+    if step < warmup_steps:
+        return base_lr * float(step + 1) / float(max(1, warmup_steps))
+    progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return min_lr + (base_lr - min_lr) * cosine

detection_tokens/pipeline.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import asdict
+
+import numpy as np
+import torch
+
+from .checkpoints import save_token_checkpoint
+from .config import PipelineConfig
+from .data import build_training_and_eval_splits, load_binary_eval_rows, load_pairs, save_pairs
+from .modeling import initialize_model_bundle
+from .scoring import dual_score, evaluate_binary_rows, evaluate_holdout
+from .training import train_single_token
+from .verbalization import save_verbalizations, verbalize_token_set
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def prepare_data(config: PipelineConfig) -> dict:
+    set_seed(config.seed)
+    config.output.output_root.mkdir(parents=True, exist_ok=True)
+    train_pairs, holdout_pairs, raid_eval_pairs, metadata = build_training_and_eval_splits(
+        config.data,
+        seed=config.seed,
+    )
+    save_pairs(config.output.splits_dir / "train_pairs.json", train_pairs)
+    save_pairs(config.output.splits_dir / "holdout_pairs.json", holdout_pairs)
+    save_pairs(config.output.splits_dir / "raid_holdout_pairs.json", raid_eval_pairs)
+    summary = {
+        "task_name": config.data.task_name,
+        "train_pairs": len(train_pairs),
+        "holdout_pairs": len(holdout_pairs),
+        "raid_holdout_pairs": len(raid_eval_pairs),
+        **metadata,
+    }
+    (config.output.output_root / "split_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    return summary
+
+
+def _load_saved_splits(config: PipelineConfig) -> tuple[list, list]:
+    train_pairs = load_pairs(config.output.splits_dir / "train_pairs.json")
+    holdout_pairs = load_pairs(config.output.splits_dir / "holdout_pairs.json")
+    return train_pairs, holdout_pairs
+
+
+def run_training(config: PipelineConfig) -> dict:
+    set_seed(config.seed)
+    train_pairs, holdout_pairs = _load_saved_splits(config)
+    bundle = initialize_model_bundle(config)
+
+    ai_artifacts, ai_ref_chosen, ai_ref_rejected = train_single_token(
+        bundle,
+        train_pairs,
+        holdout_pairs,
+        token=config.model.ai_token,
+        learning_rate=config.training.ai_learning_rate,
+        chosen_key="ai_text",
+        rejected_key="human_text",
+        config=config.training,
+        cache_dir=config.output.cache_dir,
+        seed=config.seed,
+    )
+    save_token_checkpoint(
+        token=config.model.ai_token,
+        token_id=bundle.tokenizer.convert_tokens_to_ids(config.model.ai_token),
+        embedding=ai_artifacts.embedding,
+        loss_history=ai_artifacts.loss_history,
+        path=config.output.tokens_dir / "ai_token.pt",
+    )
+
+    # Reuse the <ai> ref cache for <human> (chosen/rejected are swapped).
+    human_artifacts, _, _ = train_single_token(
+        bundle,
+        train_pairs,
+        holdout_pairs,
+        token=config.model.human_token,
+        learning_rate=config.training.human_learning_rate,
+        chosen_key="human_text",
+        rejected_key="ai_text",
+        config=config.training,
+        cache_dir=config.output.cache_dir,
+        seed=config.seed,
+        precomputed_ref_chosen=ai_ref_rejected,
+        precomputed_ref_rejected=ai_ref_chosen,
+    )
+    save_token_checkpoint(
+        token=config.model.human_token,
+        token_id=bundle.tokenizer.convert_tokens_to_ids(config.model.human_token),
+        embedding=human_artifacts.embedding,
+        loss_history=human_artifacts.loss_history,
+        path=config.output.tokens_dir / "human_token.pt",
+    )
+
+    summary = {
+        "train_pairs": len(train_pairs),
+        "holdout_pairs": len(holdout_pairs),
+        "ai_best_val_accuracy": ai_artifacts.best_val_accuracy,
+        "ai_final_val_accuracy": ai_artifacts.final_val_accuracy,
+        "human_best_val_accuracy": human_artifacts.best_val_accuracy,
+        "human_final_val_accuracy": human_artifacts.final_val_accuracy,
+    }
+    (config.output.output_root / "train_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    return summary
+
+
+def run_holdout_evaluation(config: PipelineConfig) -> dict:
+    set_seed(config.seed)
+    eval_path = config.output.splits_dir / config.evaluation.saved_pairs_filename
+    holdout_pairs = load_pairs(eval_path)
+    bundle = initialize_model_bundle(config)
+    return evaluate_holdout(bundle, holdout_pairs, config.output.evaluation_dir)
+
+
+def run_training_pipeline(config: PipelineConfig) -> dict:
+    split_summary = prepare_data(config)
+    train_summary = run_training(config)
+    evaluation = run_holdout_evaluation(config)
+    summary = {
+        "seed": config.seed,
+        "model_name": config.model.model_name,
+        "split_summary": split_summary,
+        "train_summary": train_summary,
+        "evaluation": evaluation,
+        "config": asdict(config),
+    }
+    (config.output.output_root / "pipeline_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    return summary
+
+
+def score_single_text(config: PipelineConfig, text: str) -> dict:
+    set_seed(config.seed)
+    bundle = initialize_model_bundle(config)
+    score, ai_logp, human_logp = dual_score(bundle, text)
+    return {
+        "text": text,
+        "score": score,
+        "ai_avg_logp": ai_logp,
+        "human_avg_logp": human_logp,
+    }
+
+
+def run_dataset_binary_evaluation(
+    config: PipelineConfig,
+    *,
+    split: str,
+    positive_text_types: set[str],
+    negative_text_types: set[str],
+    output_path=None,
+) -> dict:
+    set_seed(config.seed)
+    bundle = initialize_model_bundle(config)
+    rows = load_binary_eval_rows(
+        config.data,
+        split=split,
+        positive_text_types=positive_text_types,
+        negative_text_types=negative_text_types,
+    )
+    return evaluate_binary_rows(bundle, rows, output_path=output_path)
+
+
+def run_verbalizations(config: PipelineConfig) -> list[dict]:
+    set_seed(config.seed)
+    bundle = initialize_model_bundle(config)
+    results = [verbalize_token_set(bundle, token_set, config) for token_set in config.verbalization.token_sets]
+    save_verbalizations(config.verbalization.output_dir, results)
+    return results

detection_tokens/scoring.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .data import BinaryEvalRow, SourcePair
+from .modeling import ModelBundle, build_prompt, compute_average_logprob, compute_token_logprobs, encode_response
+
+
+@dataclass(slots=True)
+class ScoredText:
+    text: str
+    label: int
+    pair_id: str
+    score: float
+    ai_avg_logp: float
+    human_avg_logp: float
+
+
+def compute_auc(labels: list[int], scores: list[float]) -> float:
+    order = np.argsort(scores)
+    ranks = np.empty(len(scores), dtype=np.float64)
+    ranks[order] = np.arange(1, len(scores) + 1)
+    pos_count = sum(labels)
+    neg_count = len(labels) - pos_count
+    if pos_count == 0 or neg_count == 0:
+        return 0.5
+    pos_rank_sum = float(sum(rank for rank, label in zip(ranks, labels) if label == 1))
+    return (pos_rank_sum - pos_count * (pos_count + 1) / 2.0) / (pos_count * neg_count)
+
+
+def dual_score(bundle: ModelBundle, text: str) -> tuple[float, float, float]:
+    ai_prompt = build_prompt(bundle, bundle.config.model.ai_token)
+    human_prompt = build_prompt(bundle, bundle.config.model.human_token)
+    ai_ids, ai_prompt_len = encode_response(bundle, ai_prompt, text)
+    human_ids, human_prompt_len = encode_response(bundle, human_prompt, text)
+    ai_logp = float(compute_average_logprob(bundle, ai_ids, ai_prompt_len).item())
+    human_logp = float(compute_average_logprob(bundle, human_ids, human_prompt_len).item())
+    score_mode = bundle.config.scoring.score_mode
+    if score_mode == "avg_margin":
+        score = ai_logp - human_logp
+    elif score_mode == "soft_token_sigmoid":
+        ai_token_logps = compute_token_logprobs(bundle, ai_ids, ai_prompt_len)
+        human_token_logps = compute_token_logprobs(bundle, human_ids, human_prompt_len)
+        if ai_token_logps.shape[0] != human_token_logps.shape[0]:
+            token_count = min(ai_token_logps.shape[0], human_token_logps.shape[0])
+            ai_token_logps = ai_token_logps[:token_count]
+            human_token_logps = human_token_logps[:token_count]
+        margins = ai_token_logps - human_token_logps
+        tau = max(bundle.config.scoring.token_sigmoid_tau, 1.0e-6)
+        score = float(torch.sigmoid(margins / tau).mean().item())
+    else:
+        raise ValueError(f"Unsupported scoring mode: {score_mode}")
+    return score, ai_logp, human_logp
+
+
+def decision_threshold(bundle: ModelBundle) -> float:
+    if bundle.config.scoring.score_mode == "soft_token_sigmoid":
+        return 0.5
+    return 0.0
+
+
+def evaluate_holdout(bundle: ModelBundle, holdout_pairs: list[SourcePair], output_dir: Path) -> dict:
+    threshold = decision_threshold(bundle)
+    rows: list[ScoredText] = []
+    for pair in holdout_pairs:
+        for label, text in ((1, pair.ai_text), (0, pair.human_text)):
+            score, ai_logp, human_logp = dual_score(bundle, text)
+            rows.append(
+                ScoredText(
+                    text=text,
+                    label=label,
+                    pair_id=pair.pair_id,
+                    score=score,
+                    ai_avg_logp=ai_logp,
+                    human_avg_logp=human_logp,
+                )
+            )
+
+    labels = [row.label for row in rows]
+    scores = [row.score for row in rows]
+    tp = sum(int(row.label == 1 and row.score > threshold) for row in rows)
+    fp = sum(int(row.label == 0 and row.score > threshold) for row in rows)
+    tn = sum(int(row.label == 0 and row.score <= threshold) for row in rows)
+    fn = sum(int(row.label == 1 and row.score <= threshold) for row in rows)
+    precision = tp / max(1, tp + fp)
+    recall = tp / max(1, tp + fn)
+    f1 = 0.0 if precision + recall == 0.0 else 2.0 * precision * recall / (precision + recall)
+
+    pairwise_wins = 0
+    by_pair: dict[str, dict[int, ScoredText]] = {}
+    for row in rows:
+        by_pair.setdefault(row.pair_id, {})[row.label] = row
+    for item in by_pair.values():
+        pairwise_wins += int(item[1].score > item[0].score)
+
+    summary = {
+        "num_holdout_texts": len(rows),
+        "score_mode": bundle.config.scoring.score_mode,
+        "decision_threshold": threshold,
+        "auroc": compute_auc(labels, scores),
+        "accuracy": (tp + tn) / len(rows),
+        "pairwise_rate": pairwise_wins / len(holdout_pairs),
+        "mean_score_ai": float(np.mean([row.score for row in rows if row.label == 1])),
+        "mean_score_human": float(np.mean([row.score for row in rows if row.label == 0])),
+        "tp": tp,
+        "fp": fp,
+        "tn": tn,
+        "fn": fn,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+    }
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / "dual_eval_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    (output_dir / "dual_eval_details.json").write_text(
+        json.dumps([asdict(row) for row in rows], indent=2),
+        encoding="utf-8",
+    )
+    return summary
+
+
+def evaluate_binary_rows(
+    bundle: ModelBundle,
+    rows: list[BinaryEvalRow],
+    output_path: Path | None = None,
+) -> dict:
+    threshold = decision_threshold(bundle)
+    scores: list[float] = []
+    labels: list[int] = []
+    details = []
+    for row in rows:
+        score, ai_logp, human_logp = dual_score(bundle, row.text)
+        scores.append(score)
+        labels.append(row.label)
+        details.append(
+            {
+                "row_id": row.row_id,
+                "label": row.label,
+                "text_type": row.text_type,
+                "model": row.model,
+                "source_id": row.source_id,
+                "score": score,
+                "ai_avg_logp": ai_logp,
+                "human_avg_logp": human_logp,
+            }
+        )
+
+    labels_array = np.asarray(labels, dtype=np.int64)
+    scores_array = np.asarray(scores, dtype=np.float64)
+    predictions = (scores_array > threshold).astype(np.int64)
+
+    tp = int(((labels_array == 1) & (predictions == 1)).sum())
+    fp = int(((labels_array == 0) & (predictions == 1)).sum())
+    tn = int(((labels_array == 0) & (predictions == 0)).sum())
+    fn = int(((labels_array == 1) & (predictions == 0)).sum())
+    precision = tp / max(1, tp + fp)
+    recall = tp / max(1, tp + fn)
+    f1 = 0.0 if precision + recall == 0.0 else 2.0 * precision * recall / (precision + recall)
+
+    summary = {
+        "num_rows": len(rows),
+        "score_mode": bundle.config.scoring.score_mode,
+        "decision_threshold": threshold,
+        "positive_rows": int(labels_array.sum()),
+        "negative_rows": int((labels_array == 0).sum()),
+        "auroc": compute_auc(labels, scores),
+        "accuracy_at_zero": float((predictions == labels_array).mean()),
+        "tp": tp,
+        "fp": fp,
+        "tn": tn,
+        "fn": fn,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+        "mean_score_positive": float(scores_array[labels_array == 1].mean()),
+        "mean_score_negative": float(scores_array[labels_array == 0].mean()),
+    }
+    if output_path is not None:
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_text(
+            json.dumps({"summary": summary, "details": details}, indent=2),
+            encoding="utf-8",
+        )
+    return summary

detection_tokens/training.py

@@ -0,0 +1,430 @@
+from __future__ import annotations
+
+import gc
+import hashlib
+import random
+from collections import Counter
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+
+from .config import TrainingConfig
+from .data import SourcePair
+from .modeling import (
+    ModelBundle,
+    _model_forward,
+    build_prompt,
+    cosine_with_floor,
+    encode_response,
+)
+
+
+@dataclass(slots=True)
+class TrainedTokenArtifacts:
+    embedding: torch.Tensor
+    loss_history: list[float]
+    best_val_accuracy: float
+    final_val_accuracy: float
+    secondary_embeddings: list[torch.Tensor] = None  # per-layer embeddings (e.g. Gemma 4)
+
+
+def _ref_cache_key(model_name: str, prompt: str, pair_ids: list[str], key: str) -> str:
+    h = hashlib.sha256()
+    h.update(model_name.encode())
+    h.update(prompt.encode())
+    h.update(key.encode())
+    for pair_id in pair_ids:
+        h.update(pair_id.encode())
+    return h.hexdigest()[:16]
+
+
+def _source_balance_key(pair: SourcePair) -> str:
+    source_id = str(pair.source_id or "").strip()
+    if source_id:
+        return source_id
+    human_text = str(pair.human_text or "").strip()
+    if human_text:
+        digest = hashlib.sha256(human_text.encode("utf-8")).hexdigest()[:16]
+        return f"human_text::{digest}"
+    return f"pair::{pair.pair_id}"
+
+
+def _source_balance_weights(train_pairs: list[SourcePair]) -> torch.Tensor:
+    keys = [_source_balance_key(pair) for pair in train_pairs]
+    counts = Counter(keys)
+    if not counts:
+        return torch.ones(len(train_pairs), dtype=torch.float32)
+    scale = len(train_pairs) / len(counts)
+    weights = torch.tensor([scale / counts[key] for key in keys], dtype=torch.float32)
+    print(
+        "[source-balance] "
+        f"pairs={len(train_pairs)} sources={len(counts)} "
+        f"mean_weight={weights.mean().item():.4f} "
+        f"min_weight={weights.min().item():.4f} "
+        f"max_weight={weights.max().item():.4f}",
+        flush=True,
+    )
+    return weights
+
+
+def _pack_encoded_batch(
+    bundle: ModelBundle,
+    encoded_list: list[tuple[torch.Tensor, int]],
+) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
+    pad_id = bundle.tokenizer.pad_token_id
+    device = bundle.model.device
+    max_len = max(ids.shape[-1] for ids, _ in encoded_list)
+    input_ids = torch.full((len(encoded_list), max_len), pad_id, device=device)
+    attention_mask = torch.zeros((len(encoded_list), max_len), dtype=torch.long, device=device)
+    prompt_lens: list[int] = []
+    for batch_index, (ids, prompt_len) in enumerate(encoded_list):
+        seq_len = ids.shape[-1]
+        offset = max_len - seq_len
+        input_ids[batch_index, offset:] = ids.squeeze().to(device)
+        attention_mask[batch_index, offset:] = 1
+        prompt_lens.append(prompt_len + offset)
+    return input_ids, attention_mask, prompt_lens
+
+
+def batch_sequence_logprobs(
+    bundle: ModelBundle,
+    encoded_list: list[tuple[torch.Tensor, int]],
+) -> torch.Tensor:
+    input_ids, attention_mask, prompt_lens = _pack_encoded_batch(bundle, encoded_list)
+    logits = _model_forward(bundle, input_ids, attention_mask).logits
+    results: list[torch.Tensor] = []
+    for batch_index, prompt_len in enumerate(prompt_lens):
+        row_logits = logits[batch_index, prompt_len - 1 : -1]
+        targets = input_ids[batch_index, prompt_len:]
+        # Accumulate response likelihoods in fp32. Summing in bf16 visibly
+        # quantizes long-sequence scores and destroys fine-grained rankings.
+        row_log_probs = F.log_softmax(row_logits.float(), dim=-1)
+        token_logps = row_log_probs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
+        results.append(token_logps.sum())
+    return torch.stack(results)
+
+
+def batched_sequence_logprobs(
+    bundle: ModelBundle,
+    encoded_list: list[tuple[torch.Tensor, int]],
+    *,
+    batch_size: int,
+    desc: str,
+    checkpoint_path: Path | None = None,
+    checkpoint_every: int = 1000,
+) -> torch.Tensor:
+    n = len(encoded_list)
+    total_batches = (n + batch_size - 1) // batch_size
+    results = torch.zeros(n)
+    start_batch = 0
+
+    if checkpoint_path is not None and checkpoint_path.exists():
+        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        if ckpt["n"] == n:
+            results = ckpt["results"]
+            start_batch = ckpt["next_batch"]
+            if start_batch >= total_batches:
+                return results
+
+    for batch_idx in tqdm(range(start_batch, total_batches), desc=desc, initial=start_batch, total=total_batches):
+        start = batch_idx * batch_size
+        batch = encoded_list[start : start + batch_size]
+        results[start : start + len(batch)] = batch_sequence_logprobs(bundle, batch).detach().cpu()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        if checkpoint_path is not None and (batch_idx + 1) % checkpoint_every == 0:
+            torch.save({"n": n, "results": results, "next_batch": batch_idx + 1}, checkpoint_path)
+
+    return results
+
+
+def _load_or_compute_ref(
+    bundle: ModelBundle,
+    encoded_list: list[tuple[torch.Tensor, int]],
+    *,
+    batch_size: int,
+    desc: str,
+    cache_dir: Path,
+    cache_key: str,
+) -> torch.Tensor:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    cache_path = cache_dir / f"{cache_key}.pt"
+    ckpt_path = cache_dir / f"{cache_key}.ckpt.pt"
+    if cache_path.exists():
+        cached = torch.load(cache_path, map_location="cpu", weights_only=True)
+        if cached.shape[0] == len(encoded_list):
+            return cached
+    with torch.no_grad():
+        result = batched_sequence_logprobs(
+            bundle, encoded_list, batch_size=batch_size, desc=desc,
+            checkpoint_path=ckpt_path, checkpoint_every=1000,
+        )
+    torch.save(result, cache_path)
+    if ckpt_path.exists():
+        ckpt_path.unlink()
+    return result
+
+
+def encode_all(
+    bundle: ModelBundle,
+    pairs: list[SourcePair],
+    *,
+    prompt: str,
+    chosen_key: str,
+    rejected_key: str,
+) -> tuple[list[tuple[torch.Tensor, int]], list[tuple[torch.Tensor, int]]]:
+    chosen = []
+    rejected = []
+    for pair in tqdm(pairs, desc="encode chosen"):
+        ids, prompt_len = encode_response(bundle, prompt, getattr(pair, chosen_key))
+        chosen.append((ids.cpu(), prompt_len))
+    for pair in tqdm(pairs, desc="encode rejected"):
+        ids, prompt_len = encode_response(bundle, prompt, getattr(pair, rejected_key))
+        rejected.append((ids.cpu(), prompt_len))
+    return chosen, rejected
+
+
+@torch.no_grad()
+def eval_pair_accuracy(
+    bundle: ModelBundle,
+    pairs: list[SourcePair],
+    *,
+    prompt: str,
+    chosen_key: str,
+    rejected_key: str,
+    n: int,
+    batch_size: int,
+) -> float:
+    subset = pairs[: min(n, len(pairs))]
+    if not subset:
+        return 0.0
+    chosen_enc = []
+    rejected_enc = []
+    for pair in subset:
+        chosen_ids, chosen_prompt_len = encode_response(bundle, prompt, getattr(pair, chosen_key))
+        rejected_ids, rejected_prompt_len = encode_response(bundle, prompt, getattr(pair, rejected_key))
+        chosen_enc.append((chosen_ids.cpu(), chosen_prompt_len))
+        rejected_enc.append((rejected_ids.cpu(), rejected_prompt_len))
+
+    correct = 0
+    for start in range(0, len(subset), batch_size):
+        chosen_logps = batch_sequence_logprobs(bundle, chosen_enc[start : start + batch_size])
+        rejected_logps = batch_sequence_logprobs(bundle, rejected_enc[start : start + batch_size])
+        correct += int((chosen_logps > rejected_logps).sum().item())
+    return correct / len(subset)
+
+
+def train_single_token(
+    bundle: ModelBundle,
+    train_pairs: list[SourcePair],
+    holdout_pairs: list[SourcePair],
+    *,
+    token: str,
+    learning_rate: float,
+    chosen_key: str,
+    rejected_key: str,
+    config: TrainingConfig,
+    cache_dir: Path,
+    seed: int,
+    precomputed_ref_chosen: torch.Tensor | None = None,
+    precomputed_ref_rejected: torch.Tensor | None = None,
+) -> tuple[TrainedTokenArtifacts, torch.Tensor, torch.Tensor]:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    token_id = bundle.tokenizer.convert_tokens_to_ids(token)
+
+    # Warm-start: overwrite new token rows from a neutral existing vocab token
+    # (must happen BEFORE ref cache so ref logprobs reflect the warm start)
+    if config.neutral_word is not None:
+        nid = bundle.tokenizer.convert_tokens_to_ids(config.neutral_word)
+        unk = bundle.tokenizer.unk_token_id
+        if nid is None or nid == unk:
+            nid = bundle.tokenizer.convert_tokens_to_ids("text") or 0
+        print(f"[{token}] warm-starting from '{config.neutral_word}' (id={nid})")
+        input_emb_ws = bundle.model.get_input_embeddings()
+        with torch.no_grad():
+            input_emb_ws.weight[token_id] = input_emb_ws.weight[nid].clone()
+            for _, m in bundle.model.named_modules():
+                if (isinstance(m, torch.nn.Embedding)
+                        and m is not input_emb_ws
+                        and m.weight.shape[0] == len(bundle.tokenizer)):
+                    m.weight[token_id] = m.weight[nid].clone()
+
+    for parameter in bundle.model.parameters():
+        parameter.requires_grad = False
+    input_emb = bundle.model.get_input_embeddings()
+    input_emb.weight.requires_grad = True
+
+    # Gemma 4's auxiliary per-layer token table must be resized and mean-filled,
+    # but its new rows remain frozen. Training it diverts signal away from the
+    # tied input/output row and degrades the verbalization channel.
+    secondary_embs: list[torch.nn.Embedding] = []
+    for name, module in bundle.model.named_modules():
+        if (isinstance(module, torch.nn.Embedding)
+                and module is not input_emb
+                and module.weight.shape[0] == len(bundle.tokenizer)):
+            expected_mean = module.weight.data[:bundle.initial_tokenizer_len].mean(
+                dim=0, dtype=torch.float32
+            ).to(dtype=module.weight.dtype)
+            if not torch.equal(module.weight.data[token_id], expected_mean):
+                raise RuntimeError(
+                    f"Secondary embedding {name} token row is not mean-initialized"
+                )
+            module.weight.requires_grad = False
+            secondary_embs.append(module)
+    if secondary_embs:
+        print(
+            f"[{token}] freezing {len(secondary_embs)} secondary "
+            "embedding(s) at mean initialization"
+        )
+
+    optimizer = torch.optim.AdamW(
+        [{"params": [input_emb.weight], "lr": learning_rate,
+          "_lr_scale": 1.0}],
+        weight_decay=0.0,
+    )
+    prompt = build_prompt(bundle, token)
+
+    chosen_enc, rejected_enc = encode_all(
+        bundle,
+        train_pairs,
+        prompt=prompt,
+        chosen_key=chosen_key,
+        rejected_key=rejected_key,
+    )
+    train_ids = [pair.pair_id for pair in train_pairs]
+    pair_weights = _source_balance_weights(train_pairs) if config.source_balance_by_source else None
+    # Cache key suffix includes neutral_word so warm-start and mean-init caches
+    # never collide (neutral_word=None → empty suffix = backward-compatible)
+    ws_suffix = f"_{config.neutral_word}" if config.neutral_word else ""
+    if precomputed_ref_chosen is not None:
+        print(f"[{token}] reusing precomputed ref_chosen")
+        ref_chosen = precomputed_ref_chosen
+    else:
+        ref_chosen = _load_or_compute_ref(
+            bundle,
+            chosen_enc,
+            batch_size=config.ref_cache_batch_size,
+            desc=f"[{token}] ref chosen",
+            cache_dir=cache_dir,
+            cache_key=_ref_cache_key(bundle.config.model.model_name, prompt, train_ids, f"train_{chosen_key}{ws_suffix}"),
+        )
+    if precomputed_ref_rejected is not None:
+        print(f"[{token}] reusing precomputed ref_rejected")
+        ref_rejected = precomputed_ref_rejected
+    else:
+        ref_rejected = _load_or_compute_ref(
+            bundle,
+            rejected_enc,
+            batch_size=config.ref_cache_batch_size,
+            desc=f"[{token}] ref rejected",
+            cache_dir=cache_dir,
+            cache_key=_ref_cache_key(bundle.config.model.model_name, prompt, train_ids, f"train_{rejected_key}{ws_suffix}"),
+        )
+
+    total_steps = config.train_steps if config.train_steps is not None else len(train_pairs) // config.batch_size
+    if total_steps == 0:
+        raise ValueError(f"Not enough train pairs ({len(train_pairs)}) for batch_size={config.batch_size}.")
+
+    rng = random.Random(seed)
+    order = list(range(len(train_pairs)))
+    rng.shuffle(order)
+    total_needed = total_steps * config.batch_size
+    if total_needed > len(order):
+        order = order * (total_needed // len(order) + 1)
+    order = order[:total_needed]
+
+    best_val_accuracy = 0.0
+    final_val_accuracy = eval_pair_accuracy(
+        bundle,
+        holdout_pairs,
+        prompt=prompt,
+        chosen_key=chosen_key,
+        rejected_key=rejected_key,
+        n=config.eval_subset_size,
+        batch_size=config.ref_cache_batch_size,
+    )
+    best_val_accuracy = max(best_val_accuracy, final_val_accuracy)
+
+    loss_history: list[float] = []
+    progress = tqdm(range(1, total_steps + 1), desc=f"[{token}] train")
+    for step in progress:
+        optimizer.zero_grad(set_to_none=True)
+        batch_ids = order[(step - 1) * config.batch_size : step * config.batch_size]
+        batch_chosen = [chosen_enc[pair_index] for pair_index in batch_ids]
+        batch_rejected = [rejected_enc[pair_index] for pair_index in batch_ids]
+        bundle.model.train()
+        chosen_logps = batch_sequence_logprobs(bundle, batch_chosen)
+        rejected_logps = batch_sequence_logprobs(bundle, batch_rejected)
+        bundle.model.eval()
+        ref_c = ref_chosen[batch_ids].to(bundle.model.device)
+        ref_r = ref_rejected[batch_ids].to(bundle.model.device)
+        dpo_losses = -F.logsigmoid(config.beta * ((chosen_logps - rejected_logps) + (ref_c - ref_r)))
+        apo_losses = -F.logsigmoid(config.beta * (chosen_logps - ref_c))
+        per_example_losses = dpo_losses + config.apo_alpha * apo_losses
+        if pair_weights is not None:
+            weights = pair_weights[batch_ids].to(bundle.model.device)
+            loss = (per_example_losses * weights).mean()
+        else:
+            loss = per_example_losses.mean()
+        loss.backward()
+        step_loss = float(loss.item())
+
+        # Mask grads to token_id row only + clip just that row (avoids clip_grad_norm_ on huge tensors)
+        with torch.no_grad():
+            for emb in [input_emb] + secondary_embs:
+                if emb.weight.grad is not None:
+                    row = emb.weight.grad[token_id]
+                    row_norm = row.norm()
+                    if row_norm > 1.0:
+                        row.mul_(1.0 / row_norm)
+                    emb.weight.grad[:token_id].zero_()
+                    emb.weight.grad[token_id + 1:].zero_()
+
+        lr = cosine_with_floor(
+            step - 1,
+            total_steps,
+            learning_rate,
+            min_lr=config.min_learning_rate,
+            warmup_steps=config.warmup_steps,
+        )
+        # For secondary embeddings, scale LR proportionally
+        for group in optimizer.param_groups:
+            scale = group.get("_lr_scale", 1.0)
+            group["lr"] = lr * scale
+        optimizer.step()
+        loss_history.append(step_loss)
+        progress.set_postfix(loss=f"{step_loss:.4f}", lr=f"{lr:.2e}")
+
+        if step % config.eval_every_steps == 0 or step == total_steps:
+            final_val_accuracy = eval_pair_accuracy(
+                bundle,
+                holdout_pairs,
+                prompt=prompt,
+                chosen_key=chosen_key,
+                rejected_key=rejected_key,
+                n=config.eval_subset_size,
+                batch_size=config.ref_cache_batch_size,
+            )
+            best_val_accuracy = max(best_val_accuracy, final_val_accuracy)
+            msg = f"[{token}] step {step:>5d}  val_acc: {final_val_accuracy:.3f}"
+            progress.write(msg)
+            print(msg, flush=True)
+
+    embedding = input_emb.weight[token_id].detach().cpu()
+    secondary_embeddings = [e.weight[token_id].detach().cpu() for e in secondary_embs] if secondary_embs else None
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    return TrainedTokenArtifacts(
+        embedding=embedding,
+        loss_history=loss_history,
+        best_val_accuracy=best_val_accuracy,
+        final_val_accuracy=final_val_accuracy,
+        secondary_embeddings=secondary_embeddings,
+    ), ref_chosen, ref_rejected

detection_tokens/verbalization.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+
+from .checkpoints import load_token_checkpoint
+from .config import PipelineConfig, VerbalizationTokenSetConfig
+from .modeling import ModelBundle
+
+
+def apply_token_pair(bundle: ModelBundle, token_dir: Path) -> None:
+    input_emb = bundle.model.get_input_embeddings()
+    ai_ckpt = load_token_checkpoint(token_dir / "ai_token.pt")
+    human_ckpt = load_token_checkpoint(token_dir / "human_token.pt")
+    ai_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.ai_token)
+    human_id = bundle.tokenizer.convert_tokens_to_ids(bundle.config.model.human_token)
+    with torch.no_grad():
+        input_emb.weight[ai_id].copy_(ai_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype))
+        input_emb.weight[human_id].copy_(
+            human_ckpt.embedding.to(input_emb.weight.device, dtype=input_emb.weight.dtype)
+        )
+
+
+@torch.inference_mode()
+def sample_text(
+    bundle: ModelBundle,
+    prompt: str,
+    *,
+    max_new_tokens: int,
+    do_sample: bool,
+    temperature: float,
+    top_p: float,
+) -> str:
+    input_ids = bundle.tokenizer(prompt, return_tensors="pt", add_special_tokens=False)["input_ids"].to(
+        bundle.model.device
+    )
+    output = bundle.model.generate(
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=bundle.tokenizer.eos_token_id,
+    )
+    return bundle.tokenizer.decode(output[0, input_ids.shape[1] :], skip_special_tokens=True).strip()
+
+
+def _chat_prompt(bundle: ModelBundle, content: str) -> str:
+    return bundle.tokenizer.apply_chat_template(
+        [{"role": "user", "content": content}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+
+@torch.inference_mode()
+def verbalize_token(bundle: ModelBundle, token: str, *, max_new_tokens: int, do_sample: bool) -> str:
+    prompt = _chat_prompt(bundle, f"Describe what qualities of text would be implied by {token}.")
+    prompt += f"{token} text typically describes text that"
+    return sample_text(
+        bundle,
+        prompt,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=1.0,
+        top_p=0.95 if do_sample else 1.0,
+    )
+
+
+@torch.inference_mode()
+def verbalize_difference(bundle: ModelBundle, *, max_new_tokens: int, do_sample: bool) -> str:
+    ai_token = bundle.config.model.ai_token
+    human_token = bundle.config.model.human_token
+    prompt = _chat_prompt(
+        bundle,
+        (
+            f"What is the difference between {ai_token} text and {human_token} text? "
+            f"Refer to {ai_token} text as A-type text and {human_token} text as B-type text. "
+            "Do not discuss the literal token strings; describe the passage qualities they accompany."
+        ),
+    )
+    prompt += (
+        f"{ai_token}, which I will refer to as A-type text, and {human_token}, "
+        "B-type text, have many similarities and differences."
+    )
+    return sample_text(
+        bundle,
+        prompt,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        temperature=1.0,
+        top_p=0.95 if do_sample else 1.0,
+    )
+
+
+def verbalize_token_set(
+    bundle: ModelBundle,
+    token_set: VerbalizationTokenSetConfig,
+    config: PipelineConfig,
+) -> dict:
+    apply_token_pair(bundle, token_set.token_dir)
+    payload = {
+        "name": token_set.name,
+        "token_dir": str(token_set.token_dir),
+        "model_name": config.model.model_name,
+        "ai_token": config.model.ai_token,
+        "human_token": config.model.human_token,
+        "greedy": {
+            "ai": verbalize_token(
+                bundle,
+                config.model.ai_token,
+                max_new_tokens=config.verbalization.max_new_tokens,
+                do_sample=False,
+            ),
+            "human": verbalize_token(
+                bundle,
+                config.model.human_token,
+                max_new_tokens=config.verbalization.max_new_tokens,
+                do_sample=False,
+            ),
+            "difference": verbalize_difference(
+                bundle,
+                max_new_tokens=config.verbalization.max_new_tokens,
+                do_sample=False,
+            ),
+        },
+        "samples": [],
+    }
+    for _ in range(config.verbalization.n_samples):
+        payload["samples"].append(
+            {
+                "ai": verbalize_token(
+                    bundle,
+                    config.model.ai_token,
+                    max_new_tokens=config.verbalization.max_new_tokens,
+                    do_sample=True,
+                ),
+                "human": verbalize_token(
+                    bundle,
+                    config.model.human_token,
+                    max_new_tokens=config.verbalization.max_new_tokens,
+                    do_sample=True,
+                ),
+                "difference": verbalize_difference(
+                    bundle,
+                    max_new_tokens=config.verbalization.max_new_tokens,
+                    do_sample=True,
+                ),
+            }
+        )
+    return payload
+
+
+def save_verbalizations(output_dir: Path, results: list[dict]) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    (output_dir / "verbalizations.json").write_text(json.dumps(results, indent=2), encoding="utf-8")
+    for result in results:
+        (output_dir / f"{result['name']}.json").write_text(json.dumps(result, indent=2), encoding="utf-8")
+