Update evaluation suite

cspell/library-words.txt

@@ -8,3 +8,4 @@ xlabel
 ylabel
 frameon
 allclose
+linalg

cspell/project-words.txt

@@ -11,3 +11,4 @@ wasserstein
 bonf
 bonferroni
 aeiou
+stylometrically

pyproject.toml

@@ -11,6 +11,7 @@ dependencies = [
     "matplotlib>=3.10.8",
     "scipy>=1.17.0",
     "seaborn>=0.13.2",
+    "sentence-transformers>=5.3.0",
 ]
 
 [build-system]
@@ -101,6 +102,8 @@ module = [
     "seaborn.*",
     "matplotlib",
     "matplotlib.*",
+    "sentence_transformers",
+    "sentence_transformers.*"
 ]
 ignore_missing_imports = true
 

src/voice/__init__.py

@@ -5,12 +5,18 @@
 LLMs for stylistic fidelity.
 """
 
+from voice.comparison import (
+    make_embedding_comparison,
+    make_stylometric_comparison,
+)
 from voice.datasets import DatasetSpec, get_dataset
-from voice.stylometry import get_metrics, make_comparison
+from voice.stylometry import get_groups, get_metrics
 
 __all__: list[str] = [
     "DatasetSpec",
     "get_metrics",
+    "get_groups",
     "get_dataset",
-    "make_comparison",
+    "make_embedding_comparison",
+    "make_stylometric_comparison",
 ]

src/voice/comparison/__init__.py

@@ -0,0 +1,18 @@
+"""
+Comparison package for VOICE.
+
+This package contains:
+- Functions for comparing distributions of stylometric metrics
+"""
+
+from voice.comparison.embedding_comparison import make_embedding_comparison
+from voice.comparison.stylometric_comparison import (
+    make_stylometric_comparison,
+    stylometric_distribution,
+)
+
+__all__: list[str] = [
+    "make_stylometric_comparison",
+    "make_embedding_comparison",
+    "stylometric_distribution",
+]

src/voice/comparison/_utils.py

@@ -0,0 +1,88 @@
+"""
+Utilities for the comparison module.
+
+This module is private by convention and not part of the public API.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import numpy as np
+
+from voice.stylometry._defaults import CALIBRATION_DEFAULTS
+
+# -----------------------------------------------------------------------------
+# Calibrated percentile
+# -----------------------------------------------------------------------------
+
+
+def calibrated_percentile(value: float, dist: np.ndarray) -> float:
+    """
+    Compute the calibrated percentile (empirical CDF) of `value` under `dist`.
+
+    This returns:
+        p = (1/N) * sum_i 1 * [dist_i <= value]
+
+    :param value: Scalar value to evaluate
+    :param dist: 1D array representing a reference distribution
+    :return: Empirical CDF value in [0, 1]
+    :raises ValueError: If `dist` is empty
+    """
+    dist = np.asarray(dist, dtype=float).ravel()
+    if dist.size == 0:
+        raise ValueError("dist must be non-empty.")
+    return float(np.mean(dist <= value))
+
+
+# -----------------------------------------------------------------------------
+# Bootstrap null distribution
+# -----------------------------------------------------------------------------
+
+
+def bootstrap_null_distribution(
+    n: int,
+    score_fn: Callable[[np.ndarray, np.ndarray], float],
+    *,
+    sample_size: int = CALIBRATION_DEFAULTS.sample_size,
+    num_iterations: int = CALIBRATION_DEFAULTS.num_iterations,
+    seed: int = CALIBRATION_DEFAULTS.seed,
+) -> np.ndarray:
+    """
+    Estimate a bootstrap null distribution via random split resampling.
+
+    For each iteration, draw 2*`sample_size` unique indices from [0, n),
+    split into two non-overlapping groups of size `sample_size`, and
+    call `score_fn` on the two index arrays to produce a scalar.
+
+    The `score_fn` need not be symmetric. Each suite is responsible for
+    closing over its own precomputed data (e.g. a metric value array or
+    an embedding matrix) when constructing the callable.
+
+    :param n: Total number of examples available (typically training set size)
+    :param score_fn: Callable(idx_a, idx_b) -> float. Receives two
+        non-overlapping index arrays of length `sample_size` and returns
+        a scalar score for that split.
+    :param sample_size: Size of each resampled subset
+    :param num_iterations: Number of Monte Carlo resampling iterations
+    :param seed: RNG seed for reproducibility
+    :return: 1D array of scores of length `num_iterations`
+    :raises ValueError: If inputs are invalid
+    """
+    if sample_size <= 0:
+        raise ValueError("sample_size must be a positive integer.")
+    if num_iterations <= 0:
+        raise ValueError("num_iterations must be a positive integer.")
+    if 2 * sample_size > n:
+        raise ValueError(
+            f"sample_size must be <= half the available examples ({n / 2})."
+        )
+
+    rng = np.random.default_rng(seed)
+    out = np.empty(num_iterations, dtype=float)
+
+    for i in range(num_iterations):
+        idx = rng.choice(n, size=2 * sample_size, replace=False)
+        out[i] = score_fn(idx[:sample_size], idx[sample_size:])
+
+    return out

src/voice/comparison/embedding_comparison.py

@@ -0,0 +1,231 @@
+"""
+Functions for comparing style via neural embedding similarity.
+
+This module contains:
+    - Embedding of text examples using a pretrained SentenceTransformer
+    - Self-similarity calibration distribution (train/train resampling)
+    - Comparison of generated completions against the true corpus
+
+All calibration is performed relative to the training split by design,
+to avoid data leakage from validation/test splits.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+from voice.comparison._utils import (
+    bootstrap_null_distribution,
+    calibrated_percentile,
+)
+from voice.datasets import VoiceDataset
+from voice.datasets.dataset import Example
+from voice.stylometry._defaults import (
+    CALIBRATION_DEFAULTS,
+    COMPARISON_DEFAULTS,
+)
+
+# -----------------------------------------------------------------------------
+# Embedding extraction
+# -----------------------------------------------------------------------------
+
+
+def embedding_matrix(
+    ds: Sequence[Example],
+    model: SentenceTransformer,
+) -> np.ndarray:
+    """
+    Embed a sequence of examples using a pretrained SentenceTransformer.
+
+    Each example's `answer` field is embedded independently. The returned
+    matrix has one row per example and one column per embedding dimension,
+    with rows L2-normalised.
+
+    :param ds: Sequence of examples, each expected to have an `answer` field
+    :param model: Pretrained SentenceTransformer model
+    :return: Float32 array of shape (len(ds), embedding_dim)
+    """
+    texts = [e.answer for e in ds]
+    embeddings = model.encode(
+        texts, convert_to_numpy=True, normalize_embeddings=True
+    )
+    return embeddings.astype(np.float32)
+
+
+def mean_embedding(embeddings: np.ndarray) -> np.ndarray:
+    """
+    Compute the L2-normalised mean of a set of embeddings.
+
+    :param embeddings: Float array of shape (n, d)
+    :return: L2-normalised 1D array of shape (d,)
+    """
+    mean = embeddings.mean(axis=0)
+    norm = np.linalg.norm(mean)
+    if norm == 0.0:
+        return mean
+    return mean / norm
+
+
+def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    """
+    Compute cosine similarity between two L2-normalised vectors.
+
+    Assumes both inputs are already L2-normalised, so this reduces
+    to a dot product.
+
+    :param a: 1D array
+    :param b: 1D array of the same shape as `a`
+    :return: Cosine similarity as a float in [-1, 1]
+    """
+    return float(np.dot(a, b))
+
+
+# -----------------------------------------------------------------------------
+# Calibration: self-similarity distribution
+# -----------------------------------------------------------------------------
+
+
+def self_similarity_distribution(
+    ds: VoiceDataset,
+    model: SentenceTransformer,
+    *,
+    sample_size: int = CALIBRATION_DEFAULTS.sample_size,
+    num_iterations: int = CALIBRATION_DEFAULTS.num_iterations,
+    seed: int = CALIBRATION_DEFAULTS.seed,
+) -> np.ndarray:
+    """
+    Estimate the self-similarity distribution via resampling.
+
+    For each iteration, we draw 2*`sample_size` unique examples from the
+    training split, split them into two non-overlapping groups of size
+    `sample_size`, and compute the cosine similarity between their
+    respective mean embeddings.
+
+    :param ds: A VoiceDataset providing a `train` split
+    :param model: Pretrained SentenceTransformer model
+    :param sample_size: Size of each resampled subset
+    :param num_iterations: Number of Monte Carlo resampling iterations
+    :param seed: RNG seed for reproducibility
+    :return: 1D array of cosine similarity values of length `num_iterations`
+    """
+    ds_train = ds.train
+    n_train = len(ds_train)
+    embeddings = embedding_matrix(ds_train, model)
+
+    return bootstrap_null_distribution(
+        n_train,
+        lambda a, b: cosine_similarity(
+            mean_embedding(embeddings[a]), mean_embedding(embeddings[b])
+        ),
+        sample_size=sample_size,
+        num_iterations=num_iterations,
+        seed=seed,
+    )
+
+
+# -----------------------------------------------------------------------------
+# Comparison result container
+# -----------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingComparisonResult:
+    """
+    Result for an embedding-based style comparison.
+
+    .. attribute :: similarity
+
+        Observed cosine similarity between the mean embedding of the
+        generated completions and the mean embedding of the true corpus.
+
+    .. attribute :: percentile
+
+        Calibrated percentile under the self-similarity distribution.
+
+    .. attribute :: tail
+
+        Right-tail extremeness (1 - percentile).
+
+    .. attribute :: flag
+
+        Boolean indicating percentile >= alpha, i.e. the observed similarity
+        is not unusually low relative to within-corpus pairs (the model is
+        indistinguishable from the reference under H0: distinguishable)
+    """
+
+    similarity: float
+    percentile: float
+    tail: float
+    flag: bool
+
+    @property
+    def score(self) -> float:
+        """Alignment score in [0, 1]; 1 = perfectly indistinguishable."""
+        return self.percentile
+
+
+# -----------------------------------------------------------------------------
+# Comparison creation
+# -----------------------------------------------------------------------------
+
+
+def make_embedding_comparison(
+    completions: Sequence[Example],
+    true_ds: VoiceDataset,
+    model: SentenceTransformer,
+    *,
+    alpha: float = COMPARISON_DEFAULTS.alpha,
+) -> EmbeddingComparisonResult:
+    """
+    Compare style of `completions` to the true corpus via embedding similarity.
+
+    We compute the cosine similarity between the mean embedding of the
+    generated completions and the mean embedding of the true corpus on the
+    same split, then calibrate this against the self-similarity distribution
+    estimated from the true training split.
+
+    H0: the model is stylistically distinguishable from the reference (its
+    similarity is unusually low). flag=True rejects this null; the observed
+    similarity is not in the low tail, concluding the model is
+    indistinguishable.
+
+    :param completions: Sequence of examples to analyse
+    :param true_ds: VoiceDataset providing the true corpus
+    :param model: Pretrained SentenceTransformer model
+    :param alpha: Significance level for hypothesis testing
+    :return: EmbeddingComparisonResult
+    :raises ValueError: If completions is empty
+    """
+    if not completions:
+        raise ValueError("completions must be non-empty.")
+
+    # Split that observed completions are generated on
+    observed_split = completions[0].split
+    ds_observed_split = true_ds[observed_split]
+
+    # Observed similarity
+    gen_embeddings = embedding_matrix(completions, model)
+    ref_embeddings = embedding_matrix(ds_observed_split, model)
+
+    similarity = cosine_similarity(
+        mean_embedding(gen_embeddings),
+        mean_embedding(ref_embeddings),
+    )
+
+    # Self-similarity calibration distribution
+    self_dist = self_similarity_distribution(true_ds, model)
+
+    p = calibrated_percentile(similarity, self_dist)
+    tail = 1.0 - p
+    flag = p >= alpha
+
+    return EmbeddingComparisonResult(
+        similarity=similarity,
+        percentile=float(p),
+        tail=float(tail),
+        flag=flag,
+    )