From 469b646340fa4ab5ba33edf196831e42c644f2d4 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:59:21 -0700 Subject: [PATCH 1/2] Sweep index size and negative ratio in run_grid_search Index size is the highest-leverage knob for tuning, but run_grid_search could only sweep top_k and min_score. With subsample() available, size and ratio can be swept too, and cheaply: reshaping the index costs about a millisecond, while re-scoring costs seconds. The loop is organised around what each setting costs. Size and ratio reshape the index (cheap) and become an outer loop; top_k forces a re-score (expensive); thresholds and aggregators stay free on the cached scores. Rows now carry the requested n_positive and neg_to_pos_ratio plus the actual n_positive_actual and 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. Backward compatible: both new arguments default to None, meaning one pass using the index exactly as given. subsample() is not called at all in that case, so any index-like object keeps working and existing behaviour is unchanged. Also fixes a latent NameError: simulation.py referenced LOG without importing logging or defining it. Nothing had triggered it because the module had no logging calls until this one. Documented in the docstring, and worth calling out for reviewers: the cost is len(sizes) x len(ratios) x len(top_k) scoring passes, and the observation texts are re-encoded on every pass even though their embeddings do not depend on the index at all. Letting callers pass pre-computed sample embeddings would collapse that to roughly one encoding pass. Deliberately out of scope here. Adds 5 tests covering the default no-subsample path, both axes together and separately, clipped counts, and the per-configuration progress logging. Co-authored-by: Cursor --- src/sentinel/simulation.py | 144 +++++++++++++++++++++++++++++++------ tests/test_simulation.py | 133 ++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 21 deletions(-) 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 From c7908aa40f11664a9a17861f6d4e3cd2ae9c3f96 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:05:09 -0700 Subject: [PATCH 2/2] Prefix the grid search's index columns to stop clobbering n_positive evaluate_groups returns an "n_positive" holding the number of positive groups in the evaluation set. run_grid_search wrote the requested index size to that same key afterwards, so the group count was silently replaced. On the default path, where no size is requested, it was replaced with None - which broke the backward-compatibility guarantee the new axes were supposed to preserve, since every existing caller reading n_positive got None instead of a count. Rename the four index-describing columns to index_n_positive, index_neg_to_pos_ratio, index_n_positive_actual and index_n_negative_actual. The prefix makes the split explicit: index_ describes the index a row was scored against, everything else describes the evaluation. Renaming evaluate_groups' n_positive instead would break already-released callers of a public function. Two existing tests asserted row["n_positive"] is None on the default path, so they were locking in the very bug being reported here; they now check the prefixed columns. Adds a regression test asserting the true positive-group count survives on both the default and swept paths, which fails before this change. Reported in review by vcai4071. Co-authored-by: Cursor --- src/sentinel/simulation.py | 26 +++++++++++++----- tests/test_simulation.py | 55 ++++++++++++++++++++++++++++++-------- 2 files changed, 63 insertions(+), 18 deletions(-) diff --git a/src/sentinel/simulation.py b/src/sentinel/simulation.py index 37b9063..57150be 100644 --- a/src/sentinel/simulation.py +++ b/src/sentinel/simulation.py @@ -590,12 +590,19 @@ def run_grid_search( Returns: 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 + min_score_to_consider, aggregator)`` combination. Alongside everything + :func:`evaluate_groups` returns, each row carries ``top_k`` and four + ``index_``-prefixed columns describing the index it was scored against: the + requested ``index_n_positive`` and ``index_neg_to_pos_ratio``, and the actual + ``index_n_positive_actual`` / ``index_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. + + The prefix is not decoration. ``evaluate_groups`` already returns an + ``n_positive`` holding the number of positive *groups* in the evaluation set, + so an unprefixed index size would overwrite it and quietly turn evaluation + metadata into an index size. """ if aggregators is None: aggregators = DEFAULT_AGGREGATORS @@ -665,10 +672,15 @@ def run_grid_search( 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 + # Every index-describing column is prefixed, because + # evaluate_groups already returns an "n_positive" meaning the + # number of positive *groups* in the evaluation set. Reusing + # that name here overwrote it, silently replacing evaluation + # metadata with an index size. + row["index_n_positive"] = n_positive + row["index_neg_to_pos_ratio"] = ratio + row["index_n_positive_actual"] = n_positive_actual + row["index_n_negative_actual"] = n_negative_actual rows.append(row) return rows diff --git a/tests/test_simulation.py b/tests/test_simulation.py index ae7521c..3256c3c 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -309,8 +309,41 @@ def test_grid_search_without_index_axes_never_subsamples(): 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) + assert all(row["index_n_positive"] is None for row in rows) + assert all(row["index_neg_to_pos_ratio"] is None for row in rows) + + +def test_grid_search_keeps_evaluate_groups_n_positive(): + """The index columns must not overwrite evaluate_groups' own metadata. + + ``n_positive`` there means "how many evaluation groups are positive", which is + unrelated to the index size the sweep varies. An unprefixed index column landed + on that key and replaced a real count with the requested size - or with None on + the default path, where no size is requested at all. + """ + groups = _two_groups() # one positive group, one negative + rows = run_grid_search( + _StubSubsamplableIndex(), + groups, + top_k_values=[3], + min_score_values=[0.0], + ) + + assert all(row["n_positive"] == 1 for row in rows) + assert all(row["n_groups"] == 2 for row in rows) + + # Still true once the index axes are actually in use. + swept = run_grid_search( + _StubSubsamplableIndex(), + groups, + top_k_values=[3], + min_score_values=[0.0], + n_positive_values=[10], + neg_to_pos_ratios=[2.0], + ) + + assert all(row["n_positive"] == 1 for row in swept) + assert all(row["index_n_positive"] == 10 for row in swept) def test_grid_search_sweeps_index_size_and_ratio(): @@ -354,12 +387,12 @@ def test_grid_search_reports_requested_and_actual_counts(): 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 + by_request = {row["index_n_positive"]: row for row in rows} + assert by_request[10]["index_n_positive_actual"] == 10 + assert by_request[10]["index_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 + assert by_request[999]["index_n_positive_actual"] == 15 + assert by_request[999]["index_n_negative_actual"] == 30 def test_grid_search_one_axis_at_a_time(): @@ -369,16 +402,16 @@ def test_grid_search_one_axis_at_a_time(): 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) + assert {row["index_n_positive"] for row in size_only} == {10, 20} + assert all(row["index_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) + assert {row["index_neg_to_pos_ratio"] for row in ratio_only} == {0.5, 1.0} + assert all(row["index_n_positive"] is None for row in ratio_only) def test_grid_search_logs_progress_per_configuration(caplog):