Skip to content

Sweep index size and negative ratio in run_grid_search - #34

Open
wxiao0421 wants to merge 1 commit into
feat/index-subsamplefrom
feat/grid-search-index-axes
Open

Sweep index size and negative ratio in run_grid_search#34
wxiao0421 wants to merge 1 commit into
feat/index-subsamplefrom
feat/grid-search-index-axes

Conversation

@wxiao0421

Copy link
Copy Markdown
Contributor

Stacked on #33 (→ #32#31). Base branch is feat/index-subsample, so this diff shows only the grid-search changes. Merge the stack in order.

Summary

subsample() from #33 is the enabler; this is the payoff. Index size is the highest-leverage knob for tuning, but run_grid_search could only sweep top_k and min_score. Now it can sweep size and ratio too:

run_grid_search(
    index, groups,
    top_k_values=[3, 5, 10],
    min_score_values=[0.0, 0.1, 0.25],
    n_positive_values=[1000, 5000, 10000],   # new
    neg_to_pos_ratios=[0.2, 1.0, 5.0],       # new
    index_seed=42,                            # new
)

The loop is organised around what each setting actually costs, which is the structure the function already had — the new axes simply belong on the expensive side, as an outer loop:

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

Verified end to end against the real examples/hate_speech_model index: a 2×2×2 sweep produced 96 rows across 4 distinct index configurations, with progress logged per configuration.

Requested vs actual counts

Each row gains four keys: n_positive and neg_to_pos_ratio (what you asked for) and n_positive_actual / n_negative_actual (what you got).

The actual counts are not redundant. A request is clipped when the index is smaller than asked for, so without them two rows can look identical while describing genuinely different indices:

n_positive=10  -> n_positive_actual=10, n_negative_actual=20
n_positive=999 -> n_positive_actual=15, n_negative_actual=30   # clipped to what exists

Backward compatibility

Both new arguments default to None, meaning "one pass, use the index exactly as given." In that case subsample() is never called, so any index-like object keeps working, including test stubs that do not implement it. The pre-existing grid-search test passes untouched.

The new columns are still emitted in the default case (as None), so the result table has a consistent shape whichever way it is called.

Drive-by fix

simulation.py referenced LOG but never imported logging or defined it. That was a latent NameError — nothing had triggered it because the module contained no logging calls until this PR added one. Caught by the new progress-logging test.

A cost issue reviewers should know about

The number of scoring passes is len(n_positive_values) × len(neg_to_pos_ratios) × len(top_k_values). The 7×7 grid the internal pipeline uses, with three top_k values, is 147 full scoring passes.

Here is the subtle part: the observation texts are re-encoded on every single pass, and their embeddings do not depend on the index at all. score_groups hands raw text to calculate_rare_class_affinity, which encodes it internally each time. So most of that 147× cost is recomputing identical numbers.

The library authors already anticipated this, at sentinel_local_index.py:

"We can add it if needed, probably by just allowing the caller to pass sample embeddings instead of text."

Deliberately out of scope here. Accepting optional pre-computed sample_embeddings and caching them once in run_grid_search would cut a 147-pass sweep to roughly one encoding pass plus 147 cheap nearest-neighbour searches. That deserves its own PR. For now the cost is documented in the docstring, and progress is logged per index configuration so a long sweep cannot be mistaken for a hang.

What breaks if this is wrong

  • If the default path changed, every existing caller silently gets different results. Guarded by a test asserting subsample() is never called when the new axes are unused.
  • If subsampling happened inside the top_k loop instead of outside it, the sweep would still produce correct-looking numbers but would rebuild the index needlessly on every scoring pass, quietly throwing away the speedup this PR exists to deliver. Asserted directly: 4 configurations produce exactly 4 subsample() calls, not 8.
  • If the actual counts were wrong, a clipped grid would be misread as a genuine size sweep and lead to wrong tuning conclusions.

Test plan

5 new tests.

  • Default path never calls subsample(), and still emits the new columns as None
  • Both axes multiply out correctly; subsample() called once per (size, ratio) pair and not once per top_k
  • index_seed is forwarded to every subsample() call
  • Requested and actual counts both reported, including the clipped case
  • Either axis can be swept alone
  • Progress is logged per index configuration
  • pytest tests/ — 87 passed (82 before this PR)
  • flake8 --config=.flake8 clean on both changed files
  • docs/check_docs_sync.py reports in sync
  • Verified end to end on the real index with a genuine subsample(), not a stub

Note on CI: .github/workflows/test.yml only triggers on pull requests targeting main, so stacked PRs show no checks. Run locally against Python 3.10; the full 3.10/3.11/3.12 matrix runs once the stack retargets to main.

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 <cursoragent@cursor.com>
@wxiao0421
wxiao0421 force-pushed the feat/grid-search-index-axes branch from 21bec3a to 613c40c Compare July 29, 2026 23:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant