diff --git a/src/sentinel/simulation.py b/src/sentinel/simulation.py index 03d9f68..37b9063 100644 --- a/src/sentinel/simulation.py +++ b/src/sentinel/simulation.py @@ -56,6 +56,7 @@ rows = compare_aggregators(scored) """ +import logging from dataclasses import dataclass from typing import TYPE_CHECKING, Callable, Dict, List, Mapping, Optional, Sequence @@ -75,6 +76,8 @@ # stays lightweight (numpy-only) and does not pull in torch / transformers. from sentinel.sentinel_local_index import SentinelLocalIndex +LOG = logging.getLogger(__name__) + # A convenient name -> function map of the built-in summarize metrics, so a # caller (or a notebook) can iterate over "all aggregators" in one line. @@ -515,6 +518,21 @@ def compare_aggregators( ] +def _row_count(embeddings: object) -> Optional[int]: + """Return the number of rows in an embedding tensor, or None if unavailable. + + Args: + embeddings: A tensor-like object, or None. + + Returns: + The row count, or None when the object has no usable shape. + """ + shape = getattr(embeddings, "shape", None) + if shape is None or len(shape) == 0: + return None + return int(shape[0]) + + def run_grid_search( index: "SentinelLocalIndex", groups: Sequence[LabeledGroup], @@ -525,13 +543,33 @@ def run_grid_search( top_n: Optional[int] = None, decision_threshold: Optional[float] = None, show_progress_bar: bool = False, + n_positive_values: Optional[Sequence[int]] = None, + neg_to_pos_ratios: Optional[Sequence[float]] = None, + index_seed: Optional[int] = None, ) -> List[Dict[str, float]]: """Sweep hyperparameters and summarize metrics, returning a flat result table. - For each ``top_k`` the groups are scored once (the expensive step), then - every combination of ``min_score_values`` x ``aggregators`` is evaluated - cheaply. The returned rows are easy to turn into a ``pandas.DataFrame`` for - plotting. + The sweep is organised around what each setting costs. Index size and ratio + reshape the index (cheap, via + :meth:`~sentinel.sentinel_local_index.SentinelLocalIndex.subsample`), ``top_k`` + forces a re-scoring (expensive), and thresholds and aggregators are evaluated + for free on the cached scores:: + + for n_positive in n_positive_values: # cheap: subsample() + for ratio in neg_to_pos_ratios: # cheap: subsample() + for top_k in top_k_values: # EXPENSIVE: re-scores + for min_score in min_score_values: # cheap + for aggregator in aggregators: # cheap + + The returned rows are easy to turn into a ``pandas.DataFrame`` for plotting. + + Cost warning: the number of scoring passes is + ``len(n_positive_values) x len(neg_to_pos_ratios) x len(top_k_values)``. A 7x7 + size/ratio grid with three ``top_k`` values is 147 full passes. Note that the + observation texts are re-encoded on every pass even though their embeddings do + not depend on the index at all, so most of that cost is recomputing identical + numbers. Letting callers pass pre-computed sample embeddings would collapse it + to roughly one encoding pass; that is a separate change. Args: index: A loaded :class:`~sentinel.sentinel_local_index.SentinelLocalIndex`. @@ -543,30 +581,94 @@ def run_grid_search( decision_threshold: Cutoff for the classification family (``None`` = best-F1). show_progress_bar: Whether to show the encoder progress bar. + n_positive_values: Index sizes to try, as counts of positive examples. + ``None`` (the default) means one pass using the index exactly as given. + neg_to_pos_ratios: Negative-to-positive ratios to try. ``None`` (the + default) means one pass using the index exactly as given. + index_seed: Seed for the subsampling, so each index configuration is + reproducible across runs. Returns: - A list of metric dicts, one per ``(top_k, min_score_to_consider, - aggregator)`` combination. Each row also includes a ``top_k`` key. + A list of metric dicts, one per ``(n_positive, neg_to_pos_ratio, top_k, + min_score_to_consider, aggregator)`` combination. Each row also includes + ``top_k``, the requested ``n_positive`` and ``neg_to_pos_ratio``, and the + actual ``n_positive_actual`` / ``n_negative_actual`` counts. The actual + counts matter because a request is clipped when the index is smaller than + asked for - without them two rows can look identical while describing + different indices. """ if aggregators is None: aggregators = DEFAULT_AGGREGATORS + # None means "one pass, use the index exactly as given", which reproduces the + # behaviour from before these axes existed. + positive_options: Sequence[Optional[int]] = ( + list(n_positive_values) if n_positive_values is not None else [None] + ) + ratio_options: Sequence[Optional[float]] = ( + list(neg_to_pos_ratios) if neg_to_pos_ratios is not None else [None] + ) + total_configurations = len(positive_options) * len(ratio_options) + rows: List[Dict[str, float]] = [] - for top_k in top_k_values: - scored = score_groups( - index, groups, top_k=top_k, show_progress_bar=show_progress_bar - ) - for min_score in min_score_values: - for name, fn in aggregators.items(): - row = evaluate_groups( - scored, - fn, - aggregator_name=name, - min_score_to_consider=min_score, - top_n=top_n, - decision_threshold=decision_threshold, + configuration = 0 + for n_positive in positive_options: + for ratio in ratio_options: + configuration += 1 + + if n_positive is None and ratio is None: + # Never call subsample() in the default case, so this keeps working + # with any index-like object and stays byte-identical to before. + working_index = index + else: + working_index = index.subsample( + n_positive=n_positive, + neg_to_pos_ratio=ratio, + seed=index_seed, + ) + + n_positive_actual = _row_count( + getattr(working_index, "positive_embeddings", None) + ) + n_negative_actual = _row_count( + getattr(working_index, "negative_embeddings", None) + ) + + if total_configurations > 1: + # A long sweep is otherwise indistinguishable from a hang. + LOG.info( + "Grid search index configuration %d/%d: n_positive=%s ratio=%s " + "(actual %s positives / %s negatives)", + configuration, + total_configurations, + n_positive, + ratio, + n_positive_actual, + n_negative_actual, + ) + + for top_k in top_k_values: + scored = score_groups( + working_index, + groups, + top_k=top_k, + show_progress_bar=show_progress_bar, ) - row["top_k"] = int(top_k) - rows.append(row) + for min_score in min_score_values: + for name, fn in aggregators.items(): + row = evaluate_groups( + scored, + fn, + aggregator_name=name, + min_score_to_consider=min_score, + top_n=top_n, + decision_threshold=decision_threshold, + ) + row["top_k"] = int(top_k) + row["n_positive"] = n_positive + row["neg_to_pos_ratio"] = ratio + row["n_positive_actual"] = n_positive_actual + row["n_negative_actual"] = n_negative_actual + rows.append(row) return rows diff --git a/tests/test_simulation.py b/tests/test_simulation.py index 2275288..ae7521c 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -19,6 +19,7 @@ that mimics ``calculate_rare_class_affinity``. """ +import logging from types import SimpleNamespace import numpy as np @@ -259,3 +260,135 @@ def test_run_grid_search_rescoring_and_rows(): # Re-scoring happens once per top_k value (each scores both groups), and NOT # again for every threshold/aggregator combination: 2 top_k x 2 groups = 4. assert len(index.calls) == 2 * len(groups) + + +class _StubSubsamplableIndex(_StubIndex): + """Stub that also records subsample() calls and reports embedding row counts.""" + + def __init__(self, n_positive=100, n_negative=500): + super().__init__() + self.subsample_calls = [] + self.positive_embeddings = np.zeros((n_positive, 4)) + self.negative_embeddings = np.zeros((n_negative, 4)) + + def subsample(self, n_positive=None, neg_to_pos_ratio=None, seed=None): + self.subsample_calls.append( + {"n_positive": n_positive, "neg_to_pos_ratio": neg_to_pos_ratio, "seed": seed} + ) + available_positive = self.positive_embeddings.shape[0] + kept_positive = ( + min(n_positive, available_positive) if n_positive else available_positive + ) + kept_negative = ( + min(int(kept_positive * neg_to_pos_ratio), self.negative_embeddings.shape[0]) + if neg_to_pos_ratio + else self.negative_embeddings.shape[0] + ) + smaller = _StubSubsamplableIndex(kept_positive, kept_negative) + # Share the call log so assertions can see scoring done via the copy. + smaller.calls = self.calls + return smaller + + +def _two_groups(): + return [ + LabeledGroup(name="pos", label=1, observations=["0.9", "0.8"]), + LabeledGroup(name="neg", label=0, observations=["0.2", "0.1"]), + ] + + +def test_grid_search_without_index_axes_never_subsamples(): + """The default path must not touch the index at all. + + This is the backward-compatibility guarantee: callers passing any index-like + object keep working, and behaviour is unchanged from before these axes existed. + """ + index = _StubSubsamplableIndex() + rows = run_grid_search(index, _two_groups(), top_k_values=[3], min_score_values=[0.1]) + + assert index.subsample_calls == [] + assert len(rows) == len(DEFAULT_AGGREGATORS) + # The new columns are still present, so the table shape is consistent. + assert all(row["n_positive"] is None for row in rows) + assert all(row["neg_to_pos_ratio"] is None for row in rows) + + +def test_grid_search_sweeps_index_size_and_ratio(): + """Both new axes multiply out, and each configuration is subsampled once.""" + index = _StubSubsamplableIndex() + rows = run_grid_search( + index, + _two_groups(), + top_k_values=[3, 5], + min_score_values=[0.0], + n_positive_values=[10, 20], + neg_to_pos_ratios=[1.0, 2.0], + index_seed=42, + ) + + # 2 sizes x 2 ratios x 2 top_k x 1 threshold x 6 aggregators. + assert len(rows) == 2 * 2 * 2 * len(DEFAULT_AGGREGATORS) + # subsample() is called once per (size, ratio) pair, NOT once per top_k: the + # whole point is that reshaping the index is cheap and re-scoring is not. + assert len(index.subsample_calls) == 4 + assert all(call["seed"] == 42 for call in index.subsample_calls) + assert {(c["n_positive"], c["neg_to_pos_ratio"]) for c in index.subsample_calls} == { + (10, 1.0), (10, 2.0), (20, 1.0), (20, 2.0) + } + # Scoring happens once per (size, ratio, top_k), each covering both groups. + assert len(index.calls) == 4 * 2 * len(_two_groups()) + + +def test_grid_search_reports_requested_and_actual_counts(): + """Actual counts are emitted, because a request is clipped on a small index. + + Without them, two rows can look identical while describing different indices. + """ + index = _StubSubsamplableIndex(n_positive=15, n_negative=500) + rows = run_grid_search( + index, + _two_groups(), + top_k_values=[3], + min_score_values=[0.0], + n_positive_values=[10, 999], # 999 exceeds the 15 available + neg_to_pos_ratios=[2.0], + ) + + by_request = {row["n_positive"]: row for row in rows} + assert by_request[10]["n_positive_actual"] == 10 + assert by_request[10]["n_negative_actual"] == 20 + # Clipped to what the index actually holds, and visible in the row. + assert by_request[999]["n_positive_actual"] == 15 + assert by_request[999]["n_negative_actual"] == 30 + + +def test_grid_search_one_axis_at_a_time(): + """Either axis can be swept on its own.""" + index = _StubSubsamplableIndex() + size_only = run_grid_search( + index, _two_groups(), top_k_values=[3], min_score_values=[0.0], + n_positive_values=[10, 20], + ) + assert {row["n_positive"] for row in size_only} == {10, 20} + assert all(row["neg_to_pos_ratio"] is None for row in size_only) + + index2 = _StubSubsamplableIndex() + ratio_only = run_grid_search( + index2, _two_groups(), top_k_values=[3], min_score_values=[0.0], + neg_to_pos_ratios=[0.5, 1.0], + ) + assert {row["neg_to_pos_ratio"] for row in ratio_only} == {0.5, 1.0} + assert all(row["n_positive"] is None for row in ratio_only) + + +def test_grid_search_logs_progress_per_configuration(caplog): + """A long sweep logs progress so it cannot be mistaken for a hang.""" + caplog.set_level(logging.INFO) + index = _StubSubsamplableIndex() + run_grid_search( + index, _two_groups(), top_k_values=[3], min_score_values=[0.0], + n_positive_values=[10, 20], neg_to_pos_ratios=[1.0], + ) + + assert "index configuration 1/2" in caplog.text + assert "index configuration 2/2" in caplog.text