Lattice_Spike / scripts /eval_baselines_dgx.py
Md Saiful Islam Sajol
Claude Opus 4.8
Sync to GitHub main fd62f0f: THRIVE v6.1 weights, v6/TriviaQA pipeline, labelling tooling
ac442f8
Raw History Blame Contribute Delete
37 kB
"""Run selected token-score baselines on a HuggingFace causal LM.
Examples:
python scripts/eval_baselines_dgx.py \
--model Qwen/Qwen2.5-0.5B-Instruct \
--prompt "What is the capital of France?" \
--answer "The capital is Paris." \
--detector epr
python scripts/eval_baselines_dgx.py \
--model Qwen/Qwen2.5-3B-Instruct \
--input examples.jsonl \
--device-map auto \
--output predictions.jsonl \
--detector epr \
--detector lookback
Input JSONL rows should contain:
{"prompt": "...", "answer": "...", "labels": [0, 1, ...]}
Labels are optional. If provided, they must already align to the tokenizer's
answer tokens for AUPRC to be computed. Rows may also provide
`hallucinated_spans`, which are answer-text spans that this script converts to
token labels with tokenizer offset mappings.
Example span-labeled row:
{
"prompt": "What is the capital of France?",
"answer": "The capital is Canberra.",
"hallucinated_spans": [{"text": "Canberra."}]
}
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import re
import sys
import time
from pathlib import Path
import numpy as np
REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from baselines import (
AttentionEntropyDetector,
BaselineDetectorName,
EPRDetector,
HallucinatedSpan,
HiddenStateProbeDetector,
LookbackLensDetector,
MethodScoreAggregate,
MultiViewAttentionDetector,
ScoreSummary,
TokenEvalRow,
TokenPredictionRow,
)
from metrics import auprc
PROMPT_TEMPLATE_RAW = "raw"
PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT = "tanazzah-llama-chat"
TANAZZAH_SYSTEM_PROMPT = "Answer the question concisely and factually."
TANAZZAH_PROMPT_MAX_LENGTH = 2048
TANAZZAH_CONTEXT_TOKEN_BUDGET = 1600
TANAZZAH_CONTEXT_KEYS = (
"context",
"knowledge",
"passage",
"evidence",
"supporting_context",
)
IMMUTABLE_REVISION_RE = re.compile(r"[0-9a-fA-F]{40,64}")
def detector_choices() -> list[str]:
return [detector.value for detector in BaselineDetectorName]
def parse_args(argv=None):
parser = argparse.ArgumentParser(description="Evaluate token-score baselines on DGX/HF models.")
parser.add_argument("--model", required=True, help="HuggingFace model id or local path.")
parser.add_argument(
"--model-revision",
default=None,
help="Revision loaded for both model and tokenizer (prefer an immutable commit SHA).",
)
parser.add_argument(
"--allow-unpinned-model",
action="store_true",
help="diagnostic only: allow a missing or mutable model revision",
)
parser.add_argument("--input", help="Optional JSONL file with prompt/answer rows.")
parser.add_argument("--prompt", help="Prompt for a single-example run.")
parser.add_argument("--answer", help="Answer for a single-example run.")
parser.add_argument("--output", help="Optional JSONL path for per-example predictions.")
parser.add_argument(
"--detector",
action="append",
required=True,
choices=detector_choices(),
help="Baseline detector to run. Repeat to run multiple detectors.",
)
parser.add_argument("--top-k", type=int, default=15, help="Top-k probabilities for EPR.")
parser.add_argument("--window-size", type=int, default=8, help="Lookback Lens trailing window size.")
parser.add_argument("--hidden-layer", type=int, default=-1, help="Hidden layer index for hidden-state probe features.")
parser.add_argument("--device", default="cuda", help="Device for non-device-map loading.")
parser.add_argument("--device-map", default=None, help='Pass "auto" for sharded/Accelerate loading.')
parser.add_argument("--dtype", default="bfloat16", choices=["auto", "float32", "float16", "bfloat16"])
parser.add_argument("--attn-implementation", default="eager", help="HF attention implementation; eager is safest for attentions.")
parser.add_argument(
"--prompt-template",
default=PROMPT_TEMPLATE_RAW,
choices=[PROMPT_TEMPLATE_RAW, PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT],
help="Prompt wrapper for scoring. Use tanazzah-llama-chat to replay Tanazzah's LLaMA generation prompt.",
)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument("--require-labels", action="store_true", help="Fail if an input row has no labels or hallucinated_spans.")
parser.add_argument("--strict-labels", action="store_true", help="Fail instead of warning when labels do not align to answer tokens.")
parser.add_argument(
"--allow-token-id-model-mismatch",
action="store_true",
help="When saved generated_token_ids came from another model, re-tokenize answer text instead of failing.",
)
parser.add_argument(
"--allow-unverified-replay",
action="store_true",
help=(
"diagnostic only: permit missing/unusable saved answer IDs to be "
"retokenized and reuse saved IDs when generation-time model/"
"tokenizer or prompt provenance is missing or differs"
),
)
return parser.parse_args(argv)
def load_examples(args):
if args.input:
examples = []
with open(args.input, "r", encoding="utf-8") as f:
for line_no, line in enumerate(f, start=1):
if not line.strip():
continue
row = json.loads(line)
if "prompt" not in row or "answer" not in row:
raise ValueError(f"{args.input}:{line_no} requires prompt and answer")
examples.append(TokenEvalRow.from_json_dict(row))
return examples
if args.prompt is None or args.answer is None:
raise ValueError("provide either --input or both --prompt and --answer")
return [TokenEvalRow(prompt=args.prompt, answer=args.answer)]
def labels_from_hallucinated_spans(answer: str, tokenizer, spans):
"""Convert answer character spans into tokenizer-aligned 0/1 labels."""
try:
encoded = tokenizer(answer, add_special_tokens=False, return_offsets_mapping=True)
except (NotImplementedError, TypeError) as exc:
raise ValueError("hallucinated_spans require a fast tokenizer with offset mappings") from exc
offsets = encoded.get("offset_mapping")
if offsets is None:
raise ValueError("tokenizer did not return offset mappings for hallucinated_spans")
span_bounds = [
(span.start, span.end)
for span in (
span if isinstance(span, HallucinatedSpan) else HallucinatedSpan.from_json_dict(span, answer)
for span in spans
)
]
for start, end in span_bounds:
if start < 0 or end < start or end > len(answer):
raise ValueError(f"invalid hallucinated span bounds: start={start}, end={end}")
labels = []
for offset in offsets:
start, end = int(offset[0]), int(offset[1])
if end <= start:
labels.append(0)
continue
overlaps_span = any(max(start, span_start) < min(end, span_end) for span_start, span_end in span_bounds)
labels.append(1 if overlaps_span else 0)
return labels
def resolve_labels(example, tokenizer):
if example.labels is not None:
return [int(label) for label in example.labels]
if example.span_labels_provided:
return labels_from_hallucinated_spans(example.answer, tokenizer, example.hallucinated_spans)
return None
def prompt_add_special_tokens(prompt_template: str) -> bool:
if prompt_template in {PROMPT_TEMPLATE_RAW, PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT}:
return True
raise ValueError(f"unknown prompt template: {prompt_template}")
def prompt_max_length(prompt_template: str) -> int | None:
if prompt_template == PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT:
return TANAZZAH_PROMPT_MAX_LENGTH
if prompt_template == PROMPT_TEMPLATE_RAW:
return None
raise ValueError(f"unknown prompt template: {prompt_template}")
def tanazzah_generation_context(extra) -> str | None:
"""Mirror 02_generate_answers._extract_context exactly."""
if not isinstance(extra, dict):
return None
for key in TANAZZAH_CONTEXT_KEYS:
value = extra.get(key)
if value:
text = str(value).strip()
if text:
return text
return None
def fit_tanazzah_generation_context(context: str, tokenizer) -> str:
"""Mirror the generator's 1,600-token context-only truncation."""
encoded = tokenizer(context, add_special_tokens=False)
token_ids = encoded["input_ids"]
if hasattr(token_ids, "tolist"):
token_ids = token_ids.tolist()
if token_ids and isinstance(token_ids[0], list):
token_ids = token_ids[0]
if len(token_ids) <= TANAZZAH_CONTEXT_TOKEN_BUDGET:
return context
return tokenizer.decode(
token_ids[:TANAZZAH_CONTEXT_TOKEN_BUDGET],
skip_special_tokens=True,
)
def scoring_prompt(example: TokenEvalRow, tokenizer, prompt_template: str) -> str:
if prompt_template == PROMPT_TEMPLATE_RAW:
return example.prompt
if prompt_template != PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT:
raise ValueError(f"unknown prompt template: {prompt_template}")
system_prompt = TANAZZAH_SYSTEM_PROMPT
extra = getattr(example.metadata, "extra", None)
context = tanazzah_generation_context(extra)
if context:
system_prompt += (
f"\nContext: {fit_tanazzah_generation_context(context, tokenizer)}"
)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": example.prompt},
]
template_kwargs = {}
v5_3 = extra.get("v5_3") if isinstance(extra, dict) else None
replay = v5_3.get("replay") if isinstance(v5_3, dict) else None
if isinstance(replay, dict) and replay.get("canonical_date_string"):
template_kwargs["date_string"] = str(replay["canonical_date_string"])
try:
# Generation (02_generate_answers.py) passes enable_thinking=False for
# tokenizers that support it (Qwen3), with a TypeError fallback for the
# rest. Replay must render the *identical* prompt string, so the same
# attempt-then-fallback is mirrored here: Qwen3 with the default
# enable_thinking=True renders a different prompt than generation used,
# which would shift every token position.
try:
return tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
**template_kwargs,
)
except TypeError:
return tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
**template_kwargs,
)
except Exception as exc:
raise ValueError("tanazzah-llama-chat requires a tokenizer with apply_chat_template") from exc
def saved_answer_token_ids(
example: TokenEvalRow,
model_id: str | None = None,
model_revision: str | None = None,
allow_model_mismatch: bool = False,
allow_unverified_revision: bool = False,
require_saved: bool = False,
) -> list[int] | None:
token_ids = getattr(example.metadata, "generated_token_ids", None)
if not token_ids:
if require_saved:
raise ValueError(
"reportable dataset replay requires saved generated_token_ids "
f"for id={example.id!r}; use --allow-unverified-replay only "
"for a clearly marked diagnostic retokenization"
)
return None
generator_model = getattr(example.metadata, "generator_model", None)
if model_id and generator_model and model_id != generator_model:
if allow_model_mismatch:
if require_saved:
raise ValueError(
"saved generated_token_ids cannot be used because their "
f"generator model is {generator_model!r}, while the scoring "
f"model is {model_id!r}; add --allow-unverified-replay only "
"for a clearly marked diagnostic retokenization"
)
return None
raise ValueError(
"saved generated_token_ids were produced by "
f"{generator_model!r}, but the scoring model is {model_id!r}; "
"use the generator model for exact replay or pass "
"--allow-token-id-model-mismatch to re-tokenize answer text"
)
if model_revision is not None and not allow_unverified_revision:
extra = getattr(example.metadata, "extra", None)
generation = extra.get("generation") if isinstance(extra, dict) else None
v5_3 = extra.get("v5_3") if isinstance(extra, dict) else None
replay = v5_3.get("replay") if isinstance(v5_3, dict) else None
generated_model_revision = (
generation.get("model_revision")
if isinstance(generation, dict)
else None
)
generated_tokenizer_revision = (
generation.get("tokenizer_revision")
if isinstance(generation, dict)
else None
)
if isinstance(replay, dict) and replay.get("status") == "canonical_teacher_forced_replay":
if str(replay.get("model_id") or "") != str(model_id or ""):
raise ValueError(
"v5.3 replay model does not match the scoring model: "
f"replay={replay.get('model_id')!r}, scoring={model_id!r}"
)
if str(replay.get("model_revision") or "") != str(model_revision):
raise ValueError(
"v5.3 replay model revision does not match the scoring revision"
)
if str(replay.get("tokenizer_id") or "") != str(model_id or ""):
raise ValueError(
"v5.3 replay tokenizer does not come from the scoring model"
)
if str(replay.get("tokenizer_revision") or "") != str(model_revision):
raise ValueError(
"v5.3 replay tokenizer revision does not match the scoring revision"
)
if replay.get("answer_decode_verified") is not True:
raise ValueError("v5.3 replay has not verified answer-token decoding")
expected_ids_sha256 = replay.get("answer_token_ids_sha256")
actual_ids_sha256 = hashlib.sha256(
json.dumps(
[int(token_id) for token_id in token_ids], separators=(",", ":")
).encode("utf-8")
).hexdigest()
if str(expected_ids_sha256 or "") != actual_ids_sha256:
raise ValueError("v5.3 replay answer token IDs have changed")
return [int(token_id) for token_id in token_ids]
if str(generated_model_revision or "") != str(model_revision):
raise ValueError(
"saved generated_token_ids cannot be verified against the "
"scoring model revision: "
f"generated={generated_model_revision!r}, "
f"scoring={model_revision!r}; use --allow-unverified-replay "
"only for a clearly marked diagnostic run"
)
if str(generated_tokenizer_revision or "") != str(model_revision):
raise ValueError(
"saved generated_token_ids do not carry a matching tokenizer "
"revision: "
f"generated={generated_tokenizer_revision!r}, "
f"scoring={model_revision!r}; exact token replay is unproven. "
"Use --allow-unverified-replay only for a clearly marked "
"diagnostic run"
)
return [int(token_id) for token_id in token_ids]
def validate_saved_answer_decoding(
example: TokenEvalRow, tokenizer, token_ids: list[int]
) -> None:
"""Reject token IDs that the loaded tokenizer does not decode to the answer."""
kwargs = {
"skip_special_tokens": False,
"clean_up_tokenization_spaces": False,
}
try:
decoded = tokenizer.decode(token_ids, **kwargs)
except TypeError:
decoded = tokenizer.decode(token_ids, skip_special_tokens=False)
if decoded != example.answer:
raise ValueError(
"saved generated_token_ids do not decode to the recorded answer "
f"for id={example.id!r}; scoring this row would misalign labels"
)
def validate_generation_prompt_replay(
example: TokenEvalRow,
tokenizer,
prompt: str,
prompt_template: str,
*,
allow_unverified: bool,
) -> str:
"""Verify the effective scoring prompt token IDs when provenance exists."""
if not getattr(example.metadata, "generated_token_ids", None):
return "not_applicable_no_saved_generation"
extra = getattr(example.metadata, "extra", None)
generation = extra.get("generation") if isinstance(extra, dict) else None
v5_3 = extra.get("v5_3") if isinstance(extra, dict) else None
replay = v5_3.get("replay") if isinstance(v5_3, dict) else None
if isinstance(replay, dict) and replay.get("status") == "canonical_teacher_forced_replay":
if str(replay.get("prompt_template") or "") != prompt_template:
raise ValueError(
"v5.3 replay prompt template does not match the scoring template"
)
expected = replay.get("effective_prompt_token_ids_sha256")
expected_ids = replay.get("effective_prompt_token_ids")
verified_status = "verified_v5_3_canonical_prompt_token_ids"
actual_identity = v5_3_tokenizer_identity(tokenizer)
for field in (
"tokenizer_backend_sha256",
"chat_template_sha256",
"special_tokens_sha256",
):
if str(replay.get(field) or "") != str(actual_identity[field] or ""):
raise ValueError(
f"loaded tokenizer does not match the sealed v5.3 {field}"
)
else:
expected = (
generation.get("effective_prompt_token_ids_sha256")
if isinstance(generation, dict)
else None
)
expected_ids = None
verified_status = "verified_generation_prompt_token_ids"
if expected is None:
if allow_unverified:
return "unverified_missing_generation_prompt_hash"
raise ValueError(
"generation-time effective prompt token IDs were not recorded for "
f"id={example.id!r}; use --allow-unverified-replay only for a "
"clearly marked diagnostic run"
)
kwargs = {
"add_special_tokens": prompt_add_special_tokens(prompt_template),
}
max_length = prompt_max_length(prompt_template)
if max_length is not None:
kwargs.update({"truncation": True, "max_length": max_length})
encoded = tokenizer(prompt, **kwargs)["input_ids"]
if hasattr(encoded, "tolist"):
encoded = encoded.tolist()
if encoded and isinstance(encoded[0], list):
encoded = encoded[0]
token_ids = [int(value) for value in encoded]
if expected_ids is not None and [int(value) for value in expected_ids] != token_ids:
if allow_unverified:
return "unverified_generation_prompt_mismatch"
raise ValueError(
"scoring prompt token IDs do not match the sealed v5.3 replay "
f"for id={example.id!r}"
)
actual = hashlib.sha256(
json.dumps(token_ids, separators=(",", ":")).encode("utf-8")
).hexdigest()
if str(expected) != actual:
if allow_unverified:
return "unverified_generation_prompt_mismatch"
raise ValueError(
"scoring prompt token IDs do not match generation-time provenance "
f"for id={example.id!r}"
)
return verified_status
def v5_3_tokenizer_identity(tokenizer) -> dict[str, str | None]:
"""Content-address the tokenizer once per process for v5.3 replay."""
cached = getattr(tokenizer, "_latticespike_v5_3_identity", None)
if isinstance(cached, dict):
return cached
try:
chat_template = tokenizer.get_chat_template()
except (AttributeError, TypeError, ValueError):
# ValueError: newer transformers raises when no template is set. A
# tokenizer without one is legitimate -- --prompt-template raw needs no
# chat template -- and "absent" is a fact worth recording, not a crash.
chat_template = getattr(tokenizer, "chat_template", None)
backend = getattr(tokenizer, "backend_tokenizer", None)
backend_json = backend.to_str() if backend is not None else None
special = {
"all_special_ids": [
int(value) for value in (getattr(tokenizer, "all_special_ids", None) or [])
],
"bos_token_id": getattr(tokenizer, "bos_token_id", None),
"eos_token_id": getattr(tokenizer, "eos_token_id", None),
"pad_token_id": getattr(tokenizer, "pad_token_id", None),
}
identity = {
"tokenizer_backend_sha256": None
if backend_json is None
else hashlib.sha256(backend_json.encode("utf-8")).hexdigest(),
"chat_template_sha256": None
if not isinstance(chat_template, str)
else hashlib.sha256(chat_template.encode("utf-8")).hexdigest(),
"special_tokens_sha256": hashlib.sha256(
json.dumps(
special,
sort_keys=True,
separators=(",", ":"),
ensure_ascii=False,
).encode("utf-8")
).hexdigest(),
}
setattr(tokenizer, "_latticespike_v5_3_identity", identity)
return identity
def validate_model_revision(
revision: str | None, *, allow_unpinned: bool
) -> None:
if allow_unpinned:
return
if IMMUTABLE_REVISION_RE.fullmatch(str(revision or "")) is None:
raise ValueError(
"--model-revision must be an immutable 40-64 digit hexadecimal "
"commit hash for reportable runs; use --allow-unpinned-model only "
"for diagnostics"
)
def resolved_component_revision(component) -> str | None:
"""Best-effort resolved Hugging Face commit for run provenance."""
for holder in (
component,
getattr(component, "config", None),
getattr(component, "init_kwargs", None),
):
if holder is None:
continue
if isinstance(holder, dict):
value = holder.get("_commit_hash") or holder.get("revision")
else:
value = getattr(holder, "_commit_hash", None) or getattr(
holder, "revision", None
)
if value:
return str(value)
for attribute in ("vocab_file", "name_or_path"):
value = getattr(component, attribute, None)
if isinstance(value, str) and "snapshots" in value:
parts = value.split(os.sep)
try:
return parts[parts.index("snapshots") + 1]
except (ValueError, IndexError):
pass
return None
def torch_dtype(dtype_name):
import torch
if dtype_name == "auto":
return "auto"
if dtype_name == "float32":
return torch.float32
if dtype_name == "float16":
return torch.float16
if dtype_name == "bfloat16":
return torch.bfloat16
raise ValueError(f"unknown dtype: {dtype_name}")
def load_model_and_tokenizer(args):
from transformers import AutoModelForCausalLM, AutoTokenizer
revision = getattr(args, "model_revision", None)
tokenizer = AutoTokenizer.from_pretrained(
args.model,
revision=revision,
trust_remote_code=args.trust_remote_code,
)
if args.prompt_template == PROMPT_TEMPLATE_TANAZZAH_LLAMA_CHAT:
# The generator's final 2,048-token safety net drops from the left so
# the question and assistant marker survive pathological prompts.
tokenizer.truncation_side = "left"
model_kwargs = {
"revision": revision,
"dtype": torch_dtype(args.dtype),
"trust_remote_code": args.trust_remote_code,
"attn_implementation": args.attn_implementation,
}
if args.device_map:
model_kwargs["device_map"] = args.device_map
try:
model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
except TypeError:
# Older Transformers/model classes may reject `dtype` (which they spell
# `torch_dtype`) or `attn_implementation`. Retry the rename on its own
# first: dropping attn_implementation as well would silently fall back
# to sdpa, which returns no attention tensors, and every attention-based
# baseline would then fail much later with an unrelated-looking error.
model_kwargs["torch_dtype"] = model_kwargs.pop("dtype")
try:
model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
except TypeError:
model_kwargs.pop("attn_implementation", None)
model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
if not args.device_map and args.device:
model.to(args.device)
model.eval()
if tokenizer.pad_token is None and tokenizer.eos_token is not None:
tokenizer.pad_token = tokenizer.eos_token
# Some Transformers tokenizer classes do not retain `_commit_hash` even
# when `from_pretrained(..., revision=<immutable SHA>)` was used. The
# requested SHA is therefore the truthful resolution fallback after a
# successful pinned load; content hashes below provide an independent
# identity check for sealed v5.3 replay.
args.resolved_model_revision = resolved_component_revision(model) or revision
args.resolved_tokenizer_revision = (
resolved_component_revision(tokenizer) or revision
)
args.resolved_tokenizer_identity = v5_3_tokenizer_identity(tokenizer)
resolved = [
value
for value in (
args.resolved_model_revision,
args.resolved_tokenizer_revision,
)
if value is not None
]
if revision is not None and any(value != str(revision) for value in resolved):
raise ValueError(
"loaded model/tokenizer revision differs from --model-revision: "
f"requested={revision!r}, model={args.resolved_model_revision!r}, "
f"tokenizer={args.resolved_tokenizer_revision!r}"
)
return model, tokenizer
def summarize(method, scores, labels=None, labelled_scores=None):
arr = np.asarray(scores, dtype=float)
auprc_value = None
if labels is not None:
auprc_scores = np.asarray(labelled_scores if labelled_scores is not None else scores, dtype=float)
if len(labels) == len(auprc_scores):
auprc_value = auprc(auprc_scores, labels)
return ScoreSummary(
method=method,
n_tokens=int(arr.size),
mean=float(arr.mean()) if arr.size else float("nan"),
std=float(arr.std()) if arr.size else float("nan"),
min=float(arr.min()) if arr.size else float("nan"),
max=float(arr.max()) if arr.size else float("nan"),
p95=float(np.percentile(arr, 95)) if arr.size else float("nan"),
auprc=auprc_value,
)
def print_table(rows: list[ScoreSummary]):
method_width = max([19] + [len(row.method) for row in rows])
print(f"\n{'method':<{method_width}} tokens mean p95 max auprc")
print(f"{'-' * method_width} ------- -------- -------- -------- --------")
for row in rows:
auprc_val = row.auprc if row.auprc is not None else float("nan")
print(
f"{row.method:<{method_width}} "
f"{row.n_tokens:>7d} "
f"{row.mean:>8.4f} "
f"{row.p95:>8.4f} "
f"{row.max:>8.4f} "
f"{auprc_val:>8.4f}"
)
def aggregate_for_method(aggregates: list[MethodScoreAggregate], method: str) -> MethodScoreAggregate:
for aggregate in aggregates:
if aggregate.method == method:
return aggregate
aggregate = MethodScoreAggregate(method=method)
aggregates.append(aggregate)
return aggregate
def selected_detector_names(args) -> list[BaselineDetectorName]:
names = []
seen = set()
for detector in args.detector:
name = BaselineDetectorName(detector)
if name in seen:
continue
seen.add(name)
names.append(name)
return names
def build_detectors(args, model, tokenizer, device_for_detectors):
detectors = []
add_special_tokens = prompt_add_special_tokens(args.prompt_template)
max_length = prompt_max_length(args.prompt_template)
for name in selected_detector_names(args):
if name == BaselineDetectorName.EPR:
detectors.append(
EPRDetector(
model=model,
tokenizer=tokenizer,
top_k=args.top_k,
device=device_for_detectors,
prompt_add_special_tokens=add_special_tokens,
prompt_max_length=max_length,
)
)
elif name == BaselineDetectorName.LOOKBACK:
detectors.append(
LookbackLensDetector(
model=model,
tokenizer=tokenizer,
window_size=args.window_size,
device=device_for_detectors,
prompt_add_special_tokens=add_special_tokens,
prompt_max_length=max_length,
)
)
elif name == BaselineDetectorName.ATTENTION_ENTROPY:
detectors.append(
AttentionEntropyDetector(
model=model,
tokenizer=tokenizer,
device=device_for_detectors,
prompt_add_special_tokens=add_special_tokens,
prompt_max_length=max_length,
)
)
elif name == BaselineDetectorName.HIDDEN_STATE_PROBE:
detectors.append(
HiddenStateProbeDetector(
model=model,
tokenizer=tokenizer,
hidden_layer=args.hidden_layer,
device=device_for_detectors,
prompt_add_special_tokens=add_special_tokens,
prompt_max_length=max_length,
)
)
elif name in {BaselineDetectorName.MULTI_VIEW_ATTENTION, BaselineDetectorName.MVA_FEATURE_LOGISTIC}:
detectors.append(
MultiViewAttentionDetector(
model=model,
tokenizer=tokenizer,
device=device_for_detectors,
prompt_add_special_tokens=add_special_tokens,
prompt_max_length=max_length,
)
)
elif name == BaselineDetectorName.MULTI_VIEW_ATTENTION_TRANSFORMER_CRF:
raise ValueError("multi_view_attention_transformer_crf is trainable; use scripts/train_faithful_baselines_dgx.py")
else: # pragma: no cover - argparse enum choices prevent this.
raise ValueError(f"unsupported detector: {name}")
return detectors
def prediction_metadata(
args,
result_metadata: dict,
used_saved_answer_token_ids: bool,
prompt_replay_status: str | None = None,
) -> dict:
"""Metadata that makes a prediction row auditable after the DGX run."""
return {
**result_metadata,
"scoring_model": args.model,
"scoring_model_revision": getattr(args, "model_revision", None),
"resolved_scoring_model_revision": getattr(
args, "resolved_model_revision", None
),
"resolved_scoring_tokenizer_revision": getattr(
args, "resolved_tokenizer_revision", None
),
"prompt_template": args.prompt_template,
"prompt_add_special_tokens": prompt_add_special_tokens(args.prompt_template),
"prompt_max_length": prompt_max_length(args.prompt_template),
"used_saved_answer_token_ids": bool(used_saved_answer_token_ids),
"allow_token_id_model_mismatch": bool(args.allow_token_id_model_mismatch),
"allow_unverified_replay": bool(args.allow_unverified_replay),
"prompt_replay_status": prompt_replay_status,
"dtype": args.dtype,
"attn_implementation": args.attn_implementation,
}
def main(argv=None):
args = parse_args(argv)
validate_model_revision(
args.model_revision, allow_unpinned=args.allow_unpinned_model
)
examples = load_examples(args)
model, tokenizer = load_model_and_tokenizer(args)
device_for_detectors = None if args.device_map else args.device
detectors = build_detectors(args, model, tokenizer, device_for_detectors)
aggregates = []
prediction_rows = []
for idx, example in enumerate(examples):
prompt = scoring_prompt(example, tokenizer, args.prompt_template)
prompt_replay_status = validate_generation_prompt_replay(
example,
tokenizer,
prompt,
args.prompt_template,
allow_unverified=args.allow_unverified_replay,
)
answer = example.answer
answer_token_ids = saved_answer_token_ids(
example,
model_id=args.model,
model_revision=args.model_revision,
allow_model_mismatch=args.allow_token_id_model_mismatch,
allow_unverified_revision=args.allow_unverified_replay,
require_saved=bool(args.input) and not args.allow_unverified_replay,
)
if answer_token_ids is not None:
validate_saved_answer_decoding(example, tokenizer, answer_token_ids)
labels = resolve_labels(example, tokenizer)
if args.require_labels and labels is None:
raise ValueError(f"example {idx} requires labels or hallucinated_spans")
scored = []
for detector in detectors:
start = time.perf_counter()
result = detector.score(prompt, answer, answer_token_ids=answer_token_ids)
scored.append((result, time.perf_counter() - start))
for result, seconds in scored:
if labels is not None and len(labels) == len(result.scores):
aggregate_for_method(aggregates, result.method).extend(result.scores, labels)
elif labels is not None:
message = (
f"labels for example {idx} have length {len(labels)} "
f"but {result.method} returned {len(result.scores)} token scores"
)
if args.strict_labels:
raise ValueError(message)
print(f"warning: {message}; skipping AUPRC labels")
aggregate_for_method(aggregates, result.method).extend(result.scores, None)
else:
aggregate_for_method(aggregates, result.method).extend(result.scores, None)
prediction_rows.append(
TokenPredictionRow(
index=idx,
id=example.id,
method=result.method,
tokens=result.tokens,
scores=result.scores,
labels=labels,
seconds=seconds,
labels_aligned=labels is not None and len(labels) == len(result.scores),
metadata=prediction_metadata(
args,
result.metadata,
answer_token_ids is not None,
prompt_replay_status,
),
prompt=example.prompt,
answer=example.answer,
hallucinated_spans=example.hallucinated_spans,
example_metadata=example.metadata,
)
)
n_tokens = len(scored[0][0].tokens) if scored else 0
print(f"[{idx + 1}/{len(examples)}] scored {n_tokens} answer tokens")
table_rows = [
summarize(aggregate.method, aggregate.scores, aggregate.labels_or_none(), aggregate.labelled_scores_or_none())
for aggregate in aggregates
]
print_table(table_rows)
if args.output:
out_path = Path(args.output)
with open(out_path, "w", encoding="utf-8") as f:
for row in prediction_rows:
f.write(json.dumps(row.to_json_dict()) + "\n")
print(f"\nwrote predictions: {out_path}")
if __name__ == "__main__":
main()