Skip to content

Add index.subsample() to resize an index without re-encoding - #33

Open
wxiao0421 wants to merge 1 commit into
feat/persist-corpus-and-deterministic-loadfrom
feat/index-subsample
Open

Add index.subsample() to resize an index without re-encoding#33
wxiao0421 wants to merge 1 commit into
feat/persist-corpus-and-deterministic-loadfrom
feat/index-subsample

Conversation

@wxiao0421

@wxiao0421 wxiao0421 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Stacked on #32. Base branch is feat/persist-corpus-and-deterministic-load, so this diff shows only subsample() and its tests. Merge #32 first.

Encoding a sentence produces the same numbers regardless of which index it ends up in, so a 1,000-example index is just a 10,000-example one with 9,000 rows deleted. There is currently no way to exploit that: resizing means re-encoding the whole corpus from scratch, which makes sweeping index size — the highest-leverage knob there is — too expensive to bother with outside the internal pipeline.

subsample() copies the rows you want instead:

smaller = index.subsample(n_positive=1000, neg_to_pos_ratio=5.0, seed=42)

On the shipped examples/hate_speech_model index, building a 3×3 grid takes 9 ms against roughly 12.5 s to re-encode the same 16,100 rows. Around 1,400× faster, and the gap widens with the grid.

Design decisions worth reviewing

  • Returns a new instance and never mutates the receiver. Callers loop over sizes, so if each call shrank the original, run 2 would start from run 1's leftovers and every later result would be silently wrong. This deliberately differs from _apply_negative_ratio, which mutates in place — fine there, because it runs once inside load().
  • Positives are chosen first, since the ratio is defined relative to how many positives survive. neg_to_pos_ratio=None leaves negatives untouched.
  • Invalid sizes raise, the one departure from the existing warn-and-continue style. A grid search that quietly ignored a bad argument would emit rows describing an index nobody asked for.
  • scale_fn, encoding kwargs and model card carry over, and the sentence model is shared rather than reloaded. A dropped scale_fn would silently change scores for models like E5.
  • Reuses _take_rows() and the private-torch.Generator seeding convention from Persist the corpus and make index loading deterministic #32, so corpus alignment is handled in one place.

This is a new method: no existing signature or behaviour changes, and nothing calls it yet. The caller arrives in #34.

Test plan

15 new tests, 82 total (was 67). The load-bearing ones assert that the original index is unchanged after repeated calls, that repeated calls are independent of each other, and that corpus texts still track their own embedding rows on both sides. Also covers seeding, clipping when over-asking, carry-over of scale_fn/kwargs/model card, corpus-less indices, and ValueError on invalid sizes. flake8 output is unchanged from base and docs/check_docs_sync.py reports in sync.

No CI checks appear because .github/workflows/test.yml only triggers on pull requests targeting main. Run locally on Python 3.10; the full 3.10/3.11/3.12 matrix runs once the stack retargets.

@vcai4071 vcai4071 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

Comment on lines +423 to +434
if n_keep is None or n_keep >= available:
if n_keep is not None and n_keep > available:
LOG.info(
"Requested %d %s examples but the index only has %d - keeping all of them.",
n_keep,
label,
available,
)
# Copy the corpus list so callers cannot mutate the original through the copy.
return embeddings, (list(corpus) if corpus is not None else None)

indices = torch.randperm(available, generator=generator)[:n_keep]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

edge case: if we are doing a grid search with positive to negative ratio, if there is one run in the search where we want to keep all the positive rows, then torch.randperm(available, generator=generator) will not run, which means the stateful generator will be at a different state, and thus the negative embedding set will be different.

Encoding a sentence produces the same numbers regardless of which index it ends
up in, so a small index is just a large one with rows removed. Today there is no
way to exploit that: resizing means re-encoding the whole corpus from scratch,
which makes sweeping index size - the highest-leverage knob there is -
prohibitively expensive for anyone outside the internal pipeline.

subsample() copies the rows you want instead. On the shipped example index,
building a 3x3 grid of (n_positive, ratio) configurations takes 9 ms, against
roughly 12.5 s to re-encode the same 16,100 rows.

Design decisions worth reviewing:

- Returns a new instance and never mutates the receiver. Callers loop over sizes,
  so if each call shrank the original, run 2 would start from run 1's leftovers
  and every result after the first would be silently wrong. This differs from
  _apply_negative_ratio, which mutates in place - fine there because it runs once
  inside load().
- Positives are chosen first, because the ratio is defined relative to how many
  positives survive.
- Invalid sizes raise rather than warn-and-continue. A grid search that quietly
  ignored a bad argument would emit rows describing an index nobody asked for.
- scale_fn, encoding kwargs and model card all carry over; a dropped scale_fn
  would silently change scores for models like E5. The sentence model is shared
  rather than reloaded, since reloading it would reintroduce the cost being removed.
- Reuses the _take_rows helper and the local-torch.Generator seeding convention,
  so corpus alignment is handled in one place across the whole library.

Adds 15 tests, including that the original index is provably unchanged after
repeated calls, and that corpus texts still track their own embedding rows on
both sides of the index.

Co-authored-by: Cursor <cursoragent@cursor.com>
@wxiao0421
wxiao0421 force-pushed the feat/index-subsample branch from 839e2ee to 2287470 Compare July 29, 2026 23:49
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.

2 participants