From 4787fdfd98f0a0ecaa31bca6f223ab4abfab9c35 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:55:04 -0700 Subject: [PATCH 1/2] Add index.subsample() to resize an index without re-encoding 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 --- src/sentinel/sentinel_local_index.py | 123 +++++++++++++++++++ tests/test_sentinel_local_index.py | 169 +++++++++++++++++++++++++++ 2 files changed, 292 insertions(+) diff --git a/src/sentinel/sentinel_local_index.py b/src/sentinel/sentinel_local_index.py index 6b73bd1..2c4d2f5 100644 --- a/src/sentinel/sentinel_local_index.py +++ b/src/sentinel/sentinel_local_index.py @@ -397,6 +397,129 @@ def _apply_negative_ratio( self.negative_embeddings.shape[0], ) + def _select_subset( + self, + embeddings: torch.Tensor, + corpus: Optional[List[str]], + n_keep: Optional[int], + generator: Optional[torch.Generator], + label: str, + ) -> Tuple[torch.Tensor, Optional[List[str]]]: + """Randomly keep n_keep rows of one side of the index, corpus included. + + Args: + embeddings: The embeddings to select from. + corpus: Matching texts, or None. + n_keep: How many rows to keep. None or a value at least as large as the + available rows keeps everything. + generator: Optional seeded generator for a reproducible choice. + label: "positive" or "negative", used in log messages. + + Returns: + Tuple of (embeddings, corpus) for the kept rows. + """ + available = embeddings.shape[0] + + 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] + # Order does not affect semantic_search, but keeping the original relative + # order makes the result far easier to diff and debug. + indices = torch.sort(indices).values + LOG.info("Keeping %d %s examples out of %d", n_keep, label, available) + return _take_rows(embeddings, corpus, indices) + + def subsample( + self, + n_positive: Optional[int] = None, + neg_to_pos_ratio: Optional[float] = None, + seed: Optional[int] = None, + ) -> "SentinelLocalIndex": + """Return a smaller copy of this index, reusing the existing embeddings. + + 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. This turns + "re-encode everything" into "copy the rows you want", which is what makes + sweeping index sizes affordable. + + Corpus texts are kept aligned with the embeddings they describe, and the + sentence model is shared with the copy rather than reloaded. + + Args: + n_positive: How many positive examples to keep. None keeps all of them. + If larger than the index, everything available is kept. + neg_to_pos_ratio: Negatives to keep per kept positive. None leaves the + negatives untouched - note that shrinking the positives alone therefore + *changes* the effective ratio, which is easy to do by accident. + seed: Optional seed making the selection reproducible. Uses a private + torch.Generator, so the caller's other randomness is unaffected. + + Returns: + A new SentinelLocalIndex. This instance is never modified. + + Raises: + ValueError: If the index has no embeddings, or if n_positive or + neg_to_pos_ratio is not positive. Unlike load()'s ratio handling, which + warns and carries on, this raises: a grid search that silently ignored a + bad argument would emit result rows describing an index you did not ask for. + """ + if self.positive_embeddings is None or self.negative_embeddings is None: + raise ValueError( + "Cannot subsample an index without both positive and negative embeddings." + ) + if n_positive is not None and n_positive <= 0: + raise ValueError(f"n_positive must be positive, got {n_positive}.") + if neg_to_pos_ratio is not None and neg_to_pos_ratio <= 0: + raise ValueError( + f"neg_to_pos_ratio must be positive, got {neg_to_pos_ratio}." + ) + + generator = torch.Generator().manual_seed(seed) if seed is not None else None + + # Positives first: the ratio is defined relative to how many positives survive, + # so that count has to be settled before the negatives can be sized. + positive_embeddings, positive_corpus = self._select_subset( + self.positive_embeddings, self.positive_corpus, n_positive, generator, "positive" + ) + + n_negative: Optional[int] = None + if neg_to_pos_ratio is not None: + n_negative = int(positive_embeddings.shape[0] * neg_to_pos_ratio) + if n_negative <= 0: + LOG.info( + "Ratio %.4f against %d positives rounds to zero negatives - keeping 1, " + "since an index with no negatives cannot score anything.", + neg_to_pos_ratio, + positive_embeddings.shape[0], + ) + n_negative = 1 + + negative_embeddings, negative_corpus = self._select_subset( + self.negative_embeddings, self.negative_corpus, n_negative, generator, "negative" + ) + + # scale_fn, encoding_kwargs and model_card all carry over. A dropped scale_fn + # would silently change scores for models like E5. + return type(self)( + sentence_model=self.sentence_model, + positive_embeddings=positive_embeddings, + negative_embeddings=negative_embeddings, + scale_fn=self.scale_fn, + encoding_additional_kwargs=self.encoding_kwargs, + positive_corpus=positive_corpus, + negative_corpus=negative_corpus, + model_card=self.model_card, + ) + def calculate_rare_class_affinity( self, text_samples: List[str], diff --git a/tests/test_sentinel_local_index.py b/tests/test_sentinel_local_index.py index ebcde29..252934c 100644 --- a/tests/test_sentinel_local_index.py +++ b/tests/test_sentinel_local_index.py @@ -637,3 +637,172 @@ def test_misaligned_corpus_is_dropped_with_warning(self, caplog): assert index.negative_corpus is None assert "does not match embedding rows" in caplog.text + + +def _labelled_pair_index(n_positive=20, n_negative=40, dim=8): + """Index where BOTH sides carry the row-number-in-the-embedding trick.""" + return SentinelLocalIndex( + sentence_model=None, + positive_embeddings=torch.arange(n_positive, dtype=torch.float32).repeat(dim, 1).T, + negative_embeddings=torch.arange(n_negative, dtype=torch.float32).repeat(dim, 1).T, + positive_corpus=[f"pos {i}" for i in range(n_positive)], + negative_corpus=[f"neg {i}" for i in range(n_negative)], + scale_fn=lambda s: s * 2, + model_card={"version": "1.0"}, + ) + + +def _assert_aligned(embeddings, corpus, prefix): + """Every row must still sit beside the text it was built from.""" + for row, text in zip(embeddings, corpus): + assert text == f"{prefix} {int(row[0].item())}" + + +class TestSubsample: + """Resizing an index without re-encoding.""" + + def test_shapes_for_both_axes(self): + """Both axes resize as asked, independently and together.""" + index = _labelled_pair_index() + + only_positive = index.subsample(n_positive=5) + assert only_positive.positive_embeddings.shape[0] == 5 + assert only_positive.negative_embeddings.shape[0] == 40 # untouched + + only_ratio = index.subsample(neg_to_pos_ratio=0.5) + assert only_ratio.positive_embeddings.shape[0] == 20 # untouched + assert only_ratio.negative_embeddings.shape[0] == 10 # 20 * 0.5 + + both = index.subsample(n_positive=6, neg_to_pos_ratio=2.0) + assert both.positive_embeddings.shape[0] == 6 + assert both.negative_embeddings.shape[0] == 12 # 6 * 2.0 + + def test_original_index_is_untouched(self): + """The most important guarantee: subsample() never mutates the receiver. + + A grid search loops over sizes. If each call shrank the original, run 2 would + start from run 1's leftovers and every result after the first would be wrong. + """ + index = _labelled_pair_index() + before_positive = index.positive_embeddings.clone() + before_negative = index.negative_embeddings.clone() + before_pos_corpus = list(index.positive_corpus) + + for _ in range(3): + index.subsample(n_positive=5, neg_to_pos_ratio=1.0, seed=1) + + assert torch.equal(index.positive_embeddings, before_positive) + assert torch.equal(index.negative_embeddings, before_negative) + assert index.positive_corpus == before_pos_corpus + assert index.negative_embeddings.shape[0] == 40 + + def test_repeated_calls_are_independent(self): + """Chained calls all measure the same starting index, not each other's output.""" + index = _labelled_pair_index() + sizes = [index.subsample(n_positive=n).positive_embeddings.shape[0] for n in (10, 8, 3)] + assert sizes == [10, 8, 3] + + def test_alignment_preserved_on_both_sides(self): + """Corpus texts follow their own embedding rows through the resize.""" + index = _labelled_pair_index() + smaller = index.subsample(n_positive=7, neg_to_pos_ratio=2.0, seed=3) + + assert len(smaller.positive_corpus) == 7 + assert len(smaller.negative_corpus) == 14 + _assert_aligned(smaller.positive_embeddings, smaller.positive_corpus, "pos") + _assert_aligned(smaller.negative_embeddings, smaller.negative_corpus, "neg") + + def test_same_seed_reproducible_different_seed_not(self): + """Seeding controls the selection, as everywhere else in the library.""" + index = _labelled_pair_index(n_positive=200, n_negative=400) + + a = index.subsample(n_positive=50, neg_to_pos_ratio=1.0, seed=42) + b = index.subsample(n_positive=50, neg_to_pos_ratio=1.0, seed=42) + c = index.subsample(n_positive=50, neg_to_pos_ratio=1.0, seed=99) + + assert torch.equal(a.positive_embeddings, b.positive_embeddings) + assert a.positive_corpus == b.positive_corpus + assert not torch.equal(a.positive_embeddings, c.positive_embeddings) + + def test_requesting_more_than_available_keeps_everything(self): + """Over-asking clips rather than erroring, matching load()'s existing behaviour.""" + index = _labelled_pair_index(n_positive=20, n_negative=40) + bigger = index.subsample(n_positive=999, neg_to_pos_ratio=999.0) + + assert bigger.positive_embeddings.shape[0] == 20 + assert bigger.negative_embeddings.shape[0] == 40 + + def test_no_arguments_returns_equivalent_copy(self): + """Both arguments None yields an equal but distinct index.""" + index = _labelled_pair_index() + copy = index.subsample() + + assert copy is not index + assert torch.equal(copy.positive_embeddings, index.positive_embeddings) + assert copy.positive_corpus == index.positive_corpus + # Mutating the copy's corpus list must not reach back into the original. + copy.positive_corpus.append("extra") + assert len(index.positive_corpus) == 20 + + def test_metadata_carries_over(self): + """scale_fn, encoding kwargs, model card and the model itself come along. + + A dropped scale_fn silently changes scores for models like E5, so this is a + correctness check rather than a tidiness one. + """ + index = _labelled_pair_index() + smaller = index.subsample(n_positive=4) + + assert smaller.scale_fn is index.scale_fn + assert smaller.scale_fn(3) == 6 + assert smaller.model_card == index.model_card + assert smaller.encoding_kwargs == index.encoding_kwargs + assert smaller.encoding_kwargs["normalize_embeddings"] is True + # The model is large and read-only, so it is shared rather than reloaded. + assert smaller.sentence_model is index.sentence_model + + def test_index_without_corpus(self): + """An index carrying no corpus subsamples fine and stays corpus-free.""" + index = SentinelLocalIndex( + sentence_model=None, + positive_embeddings=torch.rand(20, 8), + negative_embeddings=torch.rand(40, 8), + ) + smaller = index.subsample(n_positive=5, neg_to_pos_ratio=2.0, seed=1) + + assert smaller.positive_embeddings.shape[0] == 5 + assert smaller.negative_embeddings.shape[0] == 10 + assert smaller.positive_corpus is None + assert smaller.negative_corpus is None + + def test_tiny_ratio_keeps_at_least_one_negative(self): + """A ratio that rounds to zero negatives is clipped to one, not left empty.""" + index = _labelled_pair_index(n_positive=20, n_negative=40) + smaller = index.subsample(n_positive=2, neg_to_pos_ratio=0.1) # 2 * 0.1 -> 0 + + assert smaller.negative_embeddings.shape[0] == 1 + + @pytest.mark.parametrize( + "kwargs,message", + [ + ({"n_positive": 0}, "n_positive must be positive"), + ({"n_positive": -5}, "n_positive must be positive"), + ({"neg_to_pos_ratio": 0}, "neg_to_pos_ratio must be positive"), + ({"neg_to_pos_ratio": -1.0}, "neg_to_pos_ratio must be positive"), + ], + ) + def test_invalid_arguments_raise(self, kwargs, message): + """Bad sizes raise instead of being quietly ignored. + + Silently ignoring them would produce grid-search rows describing an index the + caller never asked for. + """ + index = _labelled_pair_index() + with pytest.raises(ValueError, match=message): + index.subsample(**kwargs) + + def test_missing_embeddings_raise(self): + """Subsampling an index with nothing in it is an error, not a silent no-op.""" + index = SentinelLocalIndex(sentence_model=None) + with pytest.raises(ValueError, match="without both positive and negative"): + index.subsample(n_positive=1) From 6053f7a67157940471c4b5e602980ed67d00b531 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:02:35 -0700 Subject: [PATCH 2/2] Give each index side its own generator when subsampling Both sides drew from one generator, and _select_subset returns before reaching randperm when it is keeping every row. Selecting rows advances the generator, so whether the positives drew at all decided which negatives came out: at full positive size nothing was drawn, and the negatives differed from an otherwise identical call that subsampled the positives. Two grid cells meant to differ only in positive size would quietly differ in their negatives too, which defeats the point of seeding a sweep. Derive one generator per side from the caller's seed instead, so each selection depends only on that seed and its own row count. The seeds come from a root generator rather than an offset of the caller's seed, which keeps them independent without arithmetic that could collide across nearby seeds. Adds a test where both calls reach 20 negatives by different routes - 10 positives at ratio 2.0 versus 5 at ratio 4.0 - and asserts they select the same negatives. It fails on the shared generator. Reported in review by vcai4071. Co-authored-by: Cursor --- src/sentinel/sentinel_local_index.py | 54 +++++++++++++++++++++++++--- tests/test_sentinel_local_index.py | 28 +++++++++++++++ 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/src/sentinel/sentinel_local_index.py b/src/sentinel/sentinel_local_index.py index 2c4d2f5..08e5819 100644 --- a/src/sentinel/sentinel_local_index.py +++ b/src/sentinel/sentinel_local_index.py @@ -70,6 +70,38 @@ def _corpus_if_aligned( return corpus +def _split_generators( + seed: Optional[int], +) -> Tuple[Optional[torch.Generator], Optional[torch.Generator]]: + """Derive one independent generator per index side, or (None, None) when unseeded. + + A single shared generator would couple the two sides: selecting rows advances it, + so whether the positives drew at all would decide which negatives came out. Keeping + every positive draws nothing, so a grid cell at full positive size would select + different negatives from one that subsampled, and two cells meant to differ along + one axis would quietly differ along both. + + Args: + seed: The caller's seed, or None to leave selection unseeded. + + Returns: + Tuple of (positive generator, negative generator), both None when seed is None. + """ + if seed is None: + return None, None + # Draw the two seeds from a root generator rather than offsetting the caller's seed + # by hand, which keeps them independent without inventing arithmetic that could + # collide across nearby seeds. + root = torch.Generator().manual_seed(seed) + positive_seed, negative_seed = torch.randint( + high=2**62, size=(2,), generator=root + ).tolist() + return ( + torch.Generator().manual_seed(positive_seed), + torch.Generator().manual_seed(negative_seed), + ) + + def _take_rows( embeddings: torch.Tensor, corpus: Optional[List[str]], @@ -460,8 +492,10 @@ def subsample( neg_to_pos_ratio: Negatives to keep per kept positive. None leaves the negatives untouched - note that shrinking the positives alone therefore *changes* the effective ratio, which is easy to do by accident. - seed: Optional seed making the selection reproducible. Uses a private - torch.Generator, so the caller's other randomness is unaffected. + seed: Optional seed making the selection reproducible. Uses private + torch.Generators, so the caller's other randomness is unaffected. Each + side gets its own, which means the negatives chosen for a given seed and + count do not change according to whether the positives were subsampled. Returns: A new SentinelLocalIndex. This instance is never modified. @@ -483,12 +517,18 @@ def subsample( f"neg_to_pos_ratio must be positive, got {neg_to_pos_ratio}." ) - generator = torch.Generator().manual_seed(seed) if seed is not None else None + # One generator per side, so neither side's selection depends on how many draws + # the other happened to make. See _split_generators. + positive_generator, negative_generator = _split_generators(seed) # Positives first: the ratio is defined relative to how many positives survive, # so that count has to be settled before the negatives can be sized. positive_embeddings, positive_corpus = self._select_subset( - self.positive_embeddings, self.positive_corpus, n_positive, generator, "positive" + self.positive_embeddings, + self.positive_corpus, + n_positive, + positive_generator, + "positive", ) n_negative: Optional[int] = None @@ -504,7 +544,11 @@ def subsample( n_negative = 1 negative_embeddings, negative_corpus = self._select_subset( - self.negative_embeddings, self.negative_corpus, n_negative, generator, "negative" + self.negative_embeddings, + self.negative_corpus, + n_negative, + negative_generator, + "negative", ) # scale_fn, encoding_kwargs and model_card all carry over. A dropped scale_fn diff --git a/tests/test_sentinel_local_index.py b/tests/test_sentinel_local_index.py index 252934c..36e8f8b 100644 --- a/tests/test_sentinel_local_index.py +++ b/tests/test_sentinel_local_index.py @@ -724,6 +724,34 @@ def test_same_seed_reproducible_different_seed_not(self): assert a.positive_corpus == b.positive_corpus assert not torch.equal(a.positive_embeddings, c.positive_embeddings) + def test_negatives_do_not_depend_on_the_positive_code_path(self): + """The same seed and negative count must select the same negatives either way. + + Keeping every positive returns early without drawing, while subsampling them + calls randperm. On a single shared generator that difference leaves the + generator in a different state, so the negatives would silently differ between + two grid cells that were only meant to differ in positive size. + + Both calls below land on 20 negatives by different routes: 10 positives at + ratio 2.0 (no positive draw) and 5 positives at ratio 4.0 (a positive draw). + """ + index = _labelled_pair_index(n_positive=10, n_negative=200) + + kept_all_positives = index.subsample(neg_to_pos_ratio=2.0, seed=42) + subsampled_positives = index.subsample( + n_positive=5, neg_to_pos_ratio=4.0, seed=42 + ) + + assert kept_all_positives.negative_embeddings.shape[0] == 20 + assert subsampled_positives.negative_embeddings.shape[0] == 20 + assert torch.equal( + kept_all_positives.negative_embeddings, + subsampled_positives.negative_embeddings, + ) + assert ( + kept_all_positives.negative_corpus == subsampled_positives.negative_corpus + ) + def test_requesting_more_than_available_keeps_everything(self): """Over-asking clips rather than erroring, matching load()'s existing behaviour.""" index = _labelled_pair_index(n_positive=20, n_negative=40)