Blame for baselines/structs.py

"""Shared data containers for token-level hallucination baselines."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from math import isfinite
from numbers import Integral
from typing import Any, Mapping

from evaluation.protocols import PREDICTION_SCHEMA_VERSION


class BaselineDetectorName(str, Enum):
    """Supported baseline detector names for CLI/API selection."""

    EPR = "epr"
    LOOKBACK = "lookback"
    ATTENTION_ENTROPY = "attention_entropy"
    HIDDEN_STATE_PROBE = "hidden_state_probe"
    SAPLMA = "saplma"
    HAMI = "hami"
    MULTI_VIEW_ATTENTION = "multi_view_attention"
    MVA_FEATURE_LOGISTIC = "mva_feature_logistic_baseline"
    MULTI_VIEW_ATTENTION_TRANSFORMER_CRF = "multi_view_attention_transformer_crf"


def _coerce_labels(labels, name: str, allowed: set[int]) -> list[int] | None:
    if labels is None:
        return None
    allowed_str = {str(v) for v in allowed}
    text = "/".join(str(v) for v in sorted(allowed))
    values = []
    for label in labels:
        if isinstance(label, str):
            if label not in allowed_str:
                raise ValueError(f"{name} must contain only {text} values")
            values.append(int(label))
        elif isinstance(label, Integral):
            value = int(label)
            if value not in allowed:
                raise ValueError(f"{name} must contain only {text} values")
            values.append(value)
        else:
            raise ValueError(f"{name} must contain only {text} values")
    return values


def _coerce_binary_labels(labels: list[int] | None, name: str = "labels") -> list[int] | None:
    """Strict 0/1 labels — used by baseline example/prediction structs."""
    return _coerce_labels(labels, name, {0, 1})


def _coerce_token_class_labels(labels: list[int] | None, name: str = "labels") -> list[int] | None:
    """Token-eval labels: 0=grounded, 1=hallucinated, 2=uncertain.

    The uncertain class is only emitted under the "uncertain" unknown-policy;
    other policies still yield 0/1 only.
    """
    return _coerce_labels(labels, name, {0, 1, 2})


def _coerce_finite_floats(values, name: str) -> list[float]:
    coerced = []
    for value in values:
        score = float(value)
        if not isfinite(score):
            raise ValueError(f"{name} must contain only finite values")
        coerced.append(score)
    return coerced


@dataclass
class BaselineExample:
    """One prompt/answer pair for training or evaluation."""

    prompt: str
    answer: str
    labels: list[int] | None = None
    metadata: dict[str, Any] | None = None

    def __post_init__(self) -> None:
        if not isinstance(self.prompt, str):
            raise TypeError("prompt must be a string")
        if not isinstance(self.answer, str):
            raise TypeError("answer must be a string")
        if self.labels is not None and not isinstance(self.labels, list):
            raise TypeError("labels must be a list of ints or None")
        self.labels = _coerce_binary_labels(self.labels)
        if self.metadata is None:
            self.metadata = {}
        if not isinstance(self.metadata, dict):
            raise TypeError("metadata must be a dict")


@dataclass
class HallucinatedSpan:
    """A hallucinated character span inside a generated answer."""

    start: int
    end: int
    text: str | None = None
    note: str | None = None

    def __post_init__(self) -> None:
        self.start = int(self.start)
        self.end = int(self.end)
        if self.start < 0:
            raise ValueError("span start must be nonnegative")
        if self.end < self.start:
            raise ValueError("span end must be >= start")
        if self.text is not None and not isinstance(self.text, str):
            raise TypeError("span text must be a string or None")
        if self.note is not None and not isinstance(self.note, str):
            raise TypeError("span note must be a string or None")

    @classmethod
    def from_json_dict(cls, span: Mapping[str, Any], answer: str) -> "HallucinatedSpan":
        if "start" in span and "end" in span:
            start = int(span["start"])
            end = int(span["end"])
            text = span.get("text")
        elif "text" in span:
            text = span["text"]
            start = answer.find(text)
            if start < 0:
                raise ValueError(f"hallucinated span text not found in answer: {text!r}")
            end = start + len(text)
        else:
            raise ValueError("hallucinated span requires either start/end or text")

        if end > len(answer):
            raise ValueError(f"span end exceeds answer length: {end}")
        if text is None:
            text = answer[start:end]
        elif text != answer[start:end]:
            raise ValueError("hallucinated span text does not match answer bounds")
        return cls(start=start, end=end, text=text, note=span.get("note"))

    def to_json_dict(self) -> dict[str, Any]:
        row = {"start": self.start, "end": self.end}
        if self.text is not None:
            row["text"] = self.text
        if self.note is not None:
            row["note"] = self.note
        return row


@dataclass
class LabelCounts:
    """Counts of auto-labelled token classes."""

    grounded: int = 0
    hallucinated: int = 0
    unknown: int = 0

    @classmethod
    def from_labels(cls, labels: list[str] | None) -> "LabelCounts":
        counts = cls()
        for label in labels or []:
            if label == "G":
                counts.grounded += 1
            elif label == "H":
                counts.hallucinated += 1
            elif label == "U":
                counts.unknown += 1
        return counts

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any] | None) -> "LabelCounts | None":
        if row is None:
            return None
        return cls(
            grounded=int(row.get("G", 0)),
            hallucinated=int(row.get("H", 0)),
            unknown=int(row.get("U", 0)),
        )

    def to_json_dict(self) -> dict[str, int]:
        row = {}
        if self.grounded:
            row["G"] = self.grounded
        if self.hallucinated:
            row["H"] = self.hallucinated
        if self.unknown:
            row["U"] = self.unknown
        return row


@dataclass
class TokenEvalMetadata:
    """Metadata carried with token-evaluation rows."""

    source: str | None = None
    reference_answer: str | None = None
    generator_model: str | None = None
    generator_backend: str | None = None
    auto_label_method: str | None = None
    judge_model: str | None = None
    judge_rationale: str | None = None
    extra: dict[str, Any] | None = None
    generated_tokens: list[str] | None = None
    generated_token_ids: list[int] | None = None
    generated_logprobs: list[float] | None = None
    token_label_reasons: list[str] | None = None
    t_star_index: int | None = None
    original_label_counts: LabelCounts | None = None

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any] | None) -> "TokenEvalMetadata":
        if row is None:
            return cls()
        t_star_index = row.get("t_star_index")
        return cls(
            source=row.get("source"),
            reference_answer=row.get("reference_answer"),
            generator_model=row.get("generator_model"),
            generator_backend=row.get("generator_backend"),
            auto_label_method=row.get("auto_label_method"),
            judge_model=row.get("judge_model"),
            judge_rationale=row.get("judge_rationale"),
            extra=dict(row["extra"]) if isinstance(row.get("extra"), Mapping) else None,
            generated_tokens=list(row["generated_tokens"]) if row.get("generated_tokens") is not None else None,
            generated_token_ids=[int(token_id) for token_id in row["generated_token_ids"]]
            if row.get("generated_token_ids") is not None
            else None,
            generated_logprobs=[float(logprob) for logprob in row["generated_logprobs"]]
            if row.get("generated_logprobs") is not None
            else None,
            token_label_reasons=[str(reason) for reason in row["token_label_reasons"]]
            if row.get("token_label_reasons") is not None
            else None,
            t_star_index=None if t_star_index is None else int(t_star_index),
            original_label_counts=LabelCounts.from_json_dict(row.get("original_label_counts")),
        )

    def to_json_dict(self) -> dict[str, Any]:
        row = {}
        if self.source is not None:
            row["source"] = self.source
        if self.reference_answer is not None:
            row["reference_answer"] = self.reference_answer
        if self.generator_model is not None:
            row["generator_model"] = self.generator_model
        if self.generator_backend is not None:
            row["generator_backend"] = self.generator_backend
        if self.auto_label_method is not None:
            row["auto_label_method"] = self.auto_label_method
        if self.judge_model is not None:
            row["judge_model"] = self.judge_model
        if self.judge_rationale is not None:
            row["judge_rationale"] = self.judge_rationale
        if self.extra is not None:
            row["extra"] = self.extra
        if self.generated_tokens is not None:
            row["generated_tokens"] = [str(token) for token in self.generated_tokens]
        if self.generated_token_ids is not None:
            row["generated_token_ids"] = [int(token_id) for token_id in self.generated_token_ids]
        if self.generated_logprobs is not None:
            row["generated_logprobs"] = [float(logprob) for logprob in self.generated_logprobs]
        if self.token_label_reasons is not None:
            row["token_label_reasons"] = [str(reason) for reason in self.token_label_reasons]
        if self.t_star_index is not None:
            row["t_star_index"] = self.t_star_index
        if self.original_label_counts is not None:
            row["original_label_counts"] = self.original_label_counts.to_json_dict()
        return row


@dataclass
class TokenOffset:
    """Character offset for one generated token."""

    start: int
    end: int

    def __post_init__(self) -> None:
        self.start = int(self.start)
        self.end = int(self.end)
        if self.start < 0:
            raise ValueError("token offset start must be nonnegative")
        if self.end < self.start:
            raise ValueError("token offset end must be >= start")


@dataclass
class AutoLabelledRow:
    """One raw row from auto_labelled.jsonl after JSON-boundary parsing."""

    id: str
    source: str | None
    question: str
    reference_answer: str | None
    generated_answer: str
    generator_model: str | None
    generator_backend: str | None
    auto_label_method: str | None
    t_star_index: int | None
    token_labels: list[str] | None
    gen_token_offsets: list[TokenOffset] | None
    error: str | None = None
    judge_model: str | None = None
    judge_rationale: str | None = None
    extra: dict[str, Any] | None = None
    gen_tokens: list[str] | None = None
    gen_token_ids: list[int] | None = None
    gen_logprobs: list[float] | None = None
    token_label_reasons: list[str] | None = None

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any]) -> "AutoLabelledRow":
        offsets = None
        if row.get("gen_token_offsets") is not None:
            offsets = [
                TokenOffset(start=offset[0], end=offset[1])
                for offset in row["gen_token_offsets"]
            ]
        t_star_index = row.get("t_star_index")
        return cls(
            id=row["id"],
            source=row.get("source"),
            question=row["question"],
            reference_answer=row.get("reference_answer"),
            generated_answer=row.get("generated_answer") or "",
            generator_model=row.get("generator_model"),
            generator_backend=row.get("generator_backend"),
            auto_label_method=row.get("auto_label_method"),
            t_star_index=None if t_star_index is None else int(t_star_index),
            token_labels=list(row["token_labels"]) if row.get("token_labels") is not None else None,
            gen_token_offsets=offsets,
            error=row.get("error"),
            judge_model=row.get("judge_model"),
            judge_rationale=row.get("judge_rationale"),
            extra=dict(row["extra"]) if isinstance(row.get("extra"), Mapping) else None,
            gen_tokens=list(row["gen_tokens"]) if row.get("gen_tokens") is not None else None,
            gen_token_ids=[int(token_id) for token_id in row["gen_token_ids"]]
            if row.get("gen_token_ids") is not None
            else None,
            gen_logprobs=[float(logprob) for logprob in row["gen_logprobs"]]
            if row.get("gen_logprobs") is not None
            else None,
            token_label_reasons=[str(reason) for reason in row["token_label_reasons"]]
            if row.get("token_label_reasons") is not None
            else None,
        )

    def has_unknown_labels(self) -> bool:
        return self.token_labels is not None and "U" in self.token_labels


@dataclass
class TokenEvalRow:
    """One prompt/answer row for token-level baseline evaluation."""

    prompt: str
    answer: str
    id: str | None = None
    labels: list[int] | None = None
    hallucinated_spans: list[HallucinatedSpan] | None = None
    span_labels_provided: bool | None = None
    metadata: TokenEvalMetadata | None = None

    def __post_init__(self) -> None:
        if not isinstance(self.prompt, str):
            raise TypeError("prompt must be a string")
        if not isinstance(self.answer, str):
            raise TypeError("answer must be a string")
        if self.id is not None and not isinstance(self.id, str):
            raise TypeError("id must be a string or None")
        if self.labels is not None:
            if not isinstance(self.labels, list):
                raise TypeError("labels must be a list of ints or None")
            self.labels = _coerce_token_class_labels(self.labels)
        if self.span_labels_provided is None:
            self.span_labels_provided = self.hallucinated_spans is not None
        else:
            self.span_labels_provided = bool(self.span_labels_provided)
        if self.hallucinated_spans is None:
            self.hallucinated_spans = []
        self.hallucinated_spans = [
            span if isinstance(span, HallucinatedSpan) else HallucinatedSpan.from_json_dict(span, self.answer)
            for span in self.hallucinated_spans
        ]
        if self.metadata is None:
            self.metadata = TokenEvalMetadata()
        elif not isinstance(self.metadata, TokenEvalMetadata):
            self.metadata = TokenEvalMetadata.from_json_dict(self.metadata)

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any]) -> "TokenEvalRow":
        if "prompt" not in row or "answer" not in row:
            raise ValueError("evaluation rows require prompt and answer")
        return cls(
            id=row.get("id"),
            prompt=row["prompt"],
            answer=row["answer"],
            labels=row.get("labels"),
            hallucinated_spans=row.get("hallucinated_spans"),
            span_labels_provided=row.get("hallucinated_spans") is not None,
            metadata=row.get("metadata"),
        )

    def to_json_dict(self) -> dict[str, Any]:
        row = {
            "prompt": self.prompt,
            "answer": self.answer,
            "metadata": self.metadata.to_json_dict(),
        }
        if self.span_labels_provided or self.hallucinated_spans:
            row["hallucinated_spans"] = [span.to_json_dict() for span in self.hallucinated_spans]
        if self.id is not None:
            row["id"] = self.id
        if self.labels is not None:
            row["labels"] = self.labels
        return row

    def has_label_signal(self) -> bool:
        return self.labels is not None or bool(self.span_labels_provided)


@dataclass
class TokenPredictionRow:
    """One method's token-level predictions for one evaluation row.

    Prediction artifacts preserve the evaluator's full token-label vocabulary:
    0=grounded, 1=hallucinated, and 2=uncertain. Binary metric/training
    consumers are responsible for masking uncertain tokens rather than changing
    their recorded labels.
    """

    index: int
    id: str | None
    method: str
    tokens: list[str]
    scores: list[float]
    labels: list[int] | None
    seconds: float
    labels_aligned: bool
    metadata: dict[str, Any] | None = None
    prompt: str | None = None
    answer: str | None = None
    hallucinated_spans: list[HallucinatedSpan] | None = None
    example_metadata: TokenEvalMetadata | None = None
    token_ids: list[int] | None = None
    token_indices: list[int] | None = None
    score_valid_mask: list[bool] | None = None
    comparison_mask: list[bool] | None = None
    prediction_schema_version: int = PREDICTION_SCHEMA_VERSION

    def __post_init__(self) -> None:
        self.index = int(self.index)
        if self.id is not None and not isinstance(self.id, str):
            raise TypeError("id must be a string or None")
        if not isinstance(self.method, str):
            raise TypeError("method must be a string")
        self.tokens = [str(token) for token in self.tokens]
        self.scores = _coerce_finite_floats(self.scores, "scores")
        if len(self.tokens) != len(self.scores):
            raise ValueError("tokens and scores must have the same length")
        if self.labels is not None:
            self.labels = _coerce_token_class_labels(self.labels)
        self.seconds = float(self.seconds)
        self.labels_aligned = bool(self.labels_aligned)
        if self.labels_aligned:
            if self.labels is None:
                raise ValueError("labels_aligned=True requires labels")
            if len(self.labels) != len(self.scores):
                raise ValueError("aligned labels must have the same length as scores")
        if self.metadata is None:
            self.metadata = {}
        if not isinstance(self.metadata, dict):
            raise TypeError("metadata must be a dict")
        if self.prompt is not None and not isinstance(self.prompt, str):
            raise TypeError("prompt must be a string or None")
        if self.answer is not None and not isinstance(self.answer, str):
            raise TypeError("answer must be a string or None")
        if self.hallucinated_spans is None:
            self.hallucinated_spans = []
        else:
            self.hallucinated_spans = [
                self._coerce_prediction_span(span)
                for span in self.hallucinated_spans
            ]
        if self.example_metadata is None:
            self.example_metadata = TokenEvalMetadata()
        elif not isinstance(self.example_metadata, TokenEvalMetadata):
            self.example_metadata = TokenEvalMetadata.from_json_dict(self.example_metadata)
        n = len(self.scores)
        if self.token_ids is not None:
            self.token_ids = [int(token_id) for token_id in self.token_ids]
            if len(self.token_ids) != n:
                raise ValueError("token_ids must have the same length as scores")
        if self.token_indices is not None:
            self.token_indices = [int(token_index) for token_index in self.token_indices]
            if len(self.token_indices) != n:
                raise ValueError("token_indices must have the same length as scores")
            if len(set(self.token_indices)) != n:
                raise ValueError("token_indices must be unique within a prediction row")
        if self.score_valid_mask is not None:
            self.score_valid_mask = [bool(value) for value in self.score_valid_mask]
            if len(self.score_valid_mask) != n:
                raise ValueError("score_valid_mask must have the same length as scores")
        if self.comparison_mask is not None:
            self.comparison_mask = [bool(value) for value in self.comparison_mask]
            if len(self.comparison_mask) != n:
                raise ValueError("comparison_mask must have the same length as scores")
            if self.score_valid_mask is not None and any(
                compare and not valid
                for compare, valid in zip(self.comparison_mask, self.score_valid_mask)
            ):
                raise ValueError("comparison_mask cannot include invalid score slots")
        self.prediction_schema_version = int(self.prediction_schema_version)
        if self.prediction_schema_version not in {1, PREDICTION_SCHEMA_VERSION}:
            raise ValueError("unsupported prediction_schema_version")

    def _coerce_prediction_span(self, span: HallucinatedSpan | Mapping[str, Any]) -> HallucinatedSpan:
        if isinstance(span, HallucinatedSpan):
            return span
        if self.answer is not None:
            return HallucinatedSpan.from_json_dict(span, self.answer)
        if "start" not in span or "end" not in span:
            raise ValueError("prediction span text lookup requires answer context")
        return HallucinatedSpan(
            start=span["start"],
            end=span["end"],
            text=span.get("text"),
            note=span.get("note"),
        )

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any]) -> "TokenPredictionRow":
        return cls(
            index=row["index"],
            id=row.get("id"),
            method=row["method"],
            tokens=list(row["tokens"]),
            scores=list(row["scores"]),
            labels=list(row["labels"]) if row.get("labels") is not None else None,
            seconds=row.get("seconds", 0.0),
            labels_aligned=row.get("labels_aligned", False),
            metadata=row.get("metadata"),
            prompt=row.get("prompt"),
            answer=row.get("answer"),
            hallucinated_spans=row.get("hallucinated_spans"),
            example_metadata=row.get("example_metadata"),
            token_ids=list(row["token_ids"]) if row.get("token_ids") is not None else None,
            token_indices=list(row["token_indices"]) if row.get("token_indices") is not None else None,
            score_valid_mask=list(row["score_valid_mask"])
            if row.get("score_valid_mask") is not None
            else None,
            comparison_mask=list(row["comparison_mask"])
            if row.get("comparison_mask") is not None
            else None,
            prediction_schema_version=row.get("prediction_schema_version", 1),
        )

    def to_json_dict(self) -> dict[str, Any]:
        row = {
            "index": self.index,
            "id": self.id,
            "method": self.method,
            "tokens": self.tokens,
            "scores": self.scores,
            "labels": self.labels,
            "seconds": self.seconds,
            "labels_aligned": self.labels_aligned,
        }
        if self.prediction_schema_version >= PREDICTION_SCHEMA_VERSION:
            row["prediction_schema_version"] = self.prediction_schema_version
        if self.metadata:
            row["metadata"] = self.metadata
        if self.prompt is not None:
            row["prompt"] = self.prompt
        if self.answer is not None:
            row["answer"] = self.answer
        if self.hallucinated_spans:
            row["hallucinated_spans"] = [span.to_json_dict() for span in self.hallucinated_spans]
        if self.example_metadata and self.example_metadata.to_json_dict():
            row["example_metadata"] = self.example_metadata.to_json_dict()
        if self.token_ids is not None:
            row["token_ids"] = self.token_ids
        if self.token_indices is not None:
            row["token_indices"] = self.token_indices
        if self.score_valid_mask is not None:
            row["score_valid_mask"] = self.score_valid_mask
        if self.comparison_mask is not None:
            row["comparison_mask"] = self.comparison_mask
        return row


@dataclass
class ResponsePredictionRow:
    """One response-level prediction for native EPR/WEPR evaluation."""

    index: int
    id: str | None
    method: str
    score: float
    label: int | None
    seconds: float = 0.0
    score_valid: bool = True
    comparison_valid: bool = True
    metadata: dict[str, Any] | None = None
    prediction_schema_version: int = PREDICTION_SCHEMA_VERSION

    def __post_init__(self) -> None:
        self.index = int(self.index)
        if self.id is not None and not isinstance(self.id, str):
            raise TypeError("id must be a string or None")
        if not isinstance(self.method, str):
            raise TypeError("method must be a string")
        self.score = float(self.score)
        if not isfinite(self.score):
            raise ValueError("score must be finite; use score_valid=False with a finite placeholder")
        if self.label is not None:
            self.label = int(self.label)
            if self.label not in {0, 1, 2}:
                raise ValueError("label must be 0, 1, 2, or None")
        self.seconds = float(self.seconds)
        self.score_valid = bool(self.score_valid)
        self.comparison_valid = bool(self.comparison_valid)
        if self.comparison_valid and (not self.score_valid or self.label not in {0, 1}):
            raise ValueError("comparison_valid requires a valid score and binary label")
        if self.metadata is None:
            self.metadata = {}
        if not isinstance(self.metadata, dict):
            raise TypeError("metadata must be a dict")
        self.prediction_schema_version = int(self.prediction_schema_version)
        if self.prediction_schema_version != PREDICTION_SCHEMA_VERSION:
            raise ValueError("unsupported response prediction schema version")

    @classmethod
    def from_json_dict(cls, row: Mapping[str, Any]) -> "ResponsePredictionRow":
        return cls(
            index=row["index"],
            id=row.get("id"),
            method=row["method"],
            score=row["score"],
            label=row.get("label"),
            seconds=row.get("seconds", 0.0),
            score_valid=row.get("score_valid", True),
            comparison_valid=row.get("comparison_valid", True),
            metadata=row.get("metadata"),
            prediction_schema_version=row.get(
                "prediction_schema_version", PREDICTION_SCHEMA_VERSION
            ),
        )

    def to_json_dict(self) -> dict[str, Any]:
        row = {
            "prediction_schema_version": self.prediction_schema_version,
            "index": self.index,
            "id": self.id,
            "method": self.method,
            "score": self.score,
            "label": self.label,
            "seconds": self.seconds,
            "score_valid": self.score_valid,
            "comparison_valid": self.comparison_valid,
        }
        if self.metadata:
            row["metadata"] = self.metadata
        return row


@dataclass
class MethodScoreAggregate:
    """Accumulated scores and binary metric labels for one baseline method.

    All scores contribute to score-distribution statistics. Labelled metric
    arrays contain only aligned G/H labels; U labels (class 2) are deliberately
    excluded from binary AUPRC/AUROC/F1 inputs.
    """

    method: str
    scores: list[float] | None = None
    labels: list[int] | None = None
    labelled_scores: list[float] | None = None

    def __post_init__(self) -> None:
        if self.scores is None:
            self.scores = []
        if self.labels is None:
            self.labels = []
        if self.labelled_scores is None:
            self.labelled_scores = []

    def extend(
        self,
        scores: list[float],
        labels: list[int] | None,
        comparison_mask: list[bool] | None = None,
    ) -> None:
        score_values = _coerce_finite_floats(scores, "scores")
        label_values = None
        if labels is not None:
            label_values = _coerce_token_class_labels(list(labels)) or []
            if len(label_values) != len(score_values):
                raise ValueError("aggregate labels must align with scores")
        if comparison_mask is not None:
            comparison_mask = [bool(value) for value in comparison_mask]
            if len(comparison_mask) != len(score_values):
                raise ValueError("aggregate comparison mask must align with scores")
        self.scores.extend(score_values)
        if label_values is not None:
            mask = comparison_mask or [True] * len(score_values)
            for score, label, keep in zip(score_values, label_values, mask):
                if label == 2 or not keep:
                    continue
                self.labelled_scores.append(score)
                self.labels.append(label)

    def labels_or_none(self) -> list[int] | None:
        return self.labels if self.labels else None

    def labelled_scores_or_none(self) -> list[float] | None:
        return self.labelled_scores if self.labelled_scores else None


@dataclass
class ScoreSummary:
    """Aggregate statistics for one method's token scores."""

    method: str
    n_tokens: int
    mean: float
    std: float
    min: float
    max: float
    p95: float
    auprc: float | None = None


@dataclass
class TokenScoreResult:
    """Per-token hallucination-risk scores for a generated answer."""

    tokens: list[str]
    scores: list[float]
    method: str
    metadata: dict[str, Any] | None = None

    def __post_init__(self) -> None:
        self.tokens = [str(token) for token in self.tokens]
        self.scores = _coerce_finite_floats(self.scores, "scores")
        if len(self.tokens) != len(self.scores):
            raise ValueError("tokens and scores must have the same length")
        if not isinstance(self.method, str):
            raise TypeError("method must be a string")
        if self.metadata is None:
            self.metadata = {}
        if not isinstance(self.metadata, dict):
            raise TypeError("metadata must be a dict")


@dataclass
class LogisticConfig:
    """Configuration for sklearn logistic-regression heads."""

    max_iter: int = 1000
    C: float = 1.0
    class_weight: str | dict | None = None
    solver: str = "lbfgs"
    random_state: int | None = None
    penalty: str | None = "l2"