From 1bdcd0e537f6d6e949a69c52c49dab74260bc1ac Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:14:12 -0700 Subject: [PATCH 1/2] Add SentinelLocalIndex.from_texts() to build an index in one line Building an index by hand takes eight steps, two of which fail silently when skipped: omit normalize_embeddings=True and the similarity maths quietly returns wrong numbers, and omit the corpus and explanations degrade to row numbers. No crash, no warning either way. Doing those steps inside the library, where they are tested, means a caller cannot forget a step they never have to write. Also corrects the README's save-format section, which described corpus.json as optional and written only when the index has a corpus. Both became false when that file started being written unconditionally, and the file's contents are what is optional now. Co-authored-by: Cursor --- README.md | 49 ++++++++- src/sentinel/sentinel_local_index.py | 144 +++++++++++++++++++++++- tests/test_from_texts.py | 157 +++++++++++++++++++++++++++ 3 files changed, 341 insertions(+), 9 deletions(-) create mode 100644 tests/test_from_texts.py diff --git a/README.md b/README.md index 0c55e63..7126766 100644 --- a/README.md +++ b/README.md @@ -186,6 +186,38 @@ for message, ex in result.explanations.items(): ## Creating a New Index +```python +from sentinel.sentinel_local_index import SentinelLocalIndex + +index = SentinelLocalIndex.from_texts( + positive_texts=["rare class example", "critical content example"], + negative_texts=["common class example", "typical content"], + model_name="all-MiniLM-L6-v2", +) + +index.save(path="path/to/local/index", encoder_model_name_or_path="all-MiniLM-L6-v2") +``` + +`from_texts` encodes both sides, keeps the corpus so explanations can name real +sentences, and applies the correct encoding options. It optionally downsamples the +negatives for you: + +```python +index = SentinelLocalIndex.from_texts( + positive_texts=positive_texts, + negative_texts=negative_texts, + neg_to_pos_ratio=5.0, # keep 5 negatives per positive + seed=42, # ...reproducibly +) +``` + +### Advanced: building an index manually + +Use this if you need custom encoding — a different embedding backend, precomputed +vectors, or non-standard encoding arguments. Two steps here fail *silently* if you +skip them: without `normalize_embeddings=True` the similarity maths returns wrong +numbers, and without the corpus you lose explanations. Neither raises an error. + ```python import torch from sentinel.sentinel_local_index import SentinelLocalIndex @@ -343,18 +375,23 @@ Sentinel supports both local file storage and S3 storage: The storage is abstracted using `smart_open`, making it seamless to switch between storage backends. -A saved index is a directory of up to three files: +A saved index is a directory of three files: | File | Contents | |------|----------| | `sentinel_local_index_config.json` | Encoder model name, encoding kwargs, model card | | `embeddings.safetensors` | The positive and negative embedding tensors | -| `corpus.json` | The original texts behind those embeddings — **optional** | +| `corpus.json` | The original texts behind those embeddings — contents **optional** | + +`corpus.json` is what lets explanations name the matched sentence after a reload. +Without those texts, `explanations` falls back to reporting the row number of the +match instead of its text. -`corpus.json` is written whenever the index has a corpus, and it is what lets -explanations name the matched sentence after a reload. Without it, `explanations` -falls back to reporting the row number of the match instead of its text. Indices -saved before this file existed simply lack it and continue to load normally. +The file itself is always written, holding nulls when the index has no corpus, so +that it always describes the embeddings saved beside it. Saving an index without a +corpus therefore clears any corpus already at that path, rather than leaving +behind texts that describe rows which no longer exist. Indices saved before this +file existed simply lack it and continue to load normally. ## Examples To run the notebook examples diff --git a/src/sentinel/sentinel_local_index.py b/src/sentinel/sentinel_local_index.py index 08e5819..e293757 100644 --- a/src/sentinel/sentinel_local_index.py +++ b/src/sentinel/sentinel_local_index.py @@ -39,6 +39,33 @@ LOG = logging.getLogger(__name__) +# Encoding options applied unless the caller overrides them. Normalization is not +# optional in practice: the similarity maths assumes unit vectors, and getting it +# wrong produces quietly incorrect scores rather than an error. Defined once so the +# constructor and from_texts() cannot drift apart. +DEFAULT_ENCODING_KWARGS: Mapping[str, Any] = { + "normalize_embeddings": True, +} + + +def _validate_texts(name: str, texts: Any) -> None: + """Reject the two text-argument mistakes that would otherwise fail quietly. + + Args: + name: Argument name, used in the error message. + texts: The value supplied by the caller. + + Raises: + ValueError: If a bare string was passed instead of a list, or the list is empty. + """ + if isinstance(texts, str): + raise ValueError( + f"{name} must be a list of strings, not a single string. A bare string is " + f"iterable, so it would be encoded one character at a time instead of failing." + ) + if texts is None or len(texts) == 0: + raise ValueError(f"{name} must not be empty.") + def _corpus_if_aligned( name: str, corpus: Optional[List[str]], embeddings: Optional[torch.Tensor] @@ -214,9 +241,7 @@ def __init__( else: self.negative_embeddings = torch.tensor(negative_embeddings) - self.encoding_kwargs = { - "normalize_embeddings": True, - } + self.encoding_kwargs = dict(DEFAULT_ENCODING_KWARGS) self.encoding_kwargs.update(encoding_additional_kwargs) self.positive_corpus = positive_corpus self.negative_corpus = negative_corpus @@ -271,6 +296,119 @@ def save( # Return the config for informational purposes return config + @classmethod + def from_texts( + cls, + positive_texts: List[str], + negative_texts: List[str], + model_name: str = "all-MiniLM-L6-v2", + neg_to_pos_ratio: Optional[float] = None, + batch_size: int = 256, + seed: Optional[int] = None, + encoding_additional_kwargs: Optional[Mapping[str, Any]] = None, + model_card: Optional[Mapping[str, Any]] = None, + show_progress_bar: bool = False, + ) -> "SentinelLocalIndex": + """Build an index from raw text in one call. + + Doing this by hand takes eight steps, two of which fail *silently* when + skipped: omit ``normalize_embeddings=True`` and the similarity maths quietly + returns wrong numbers, and omit the corpus and you lose explanations. No + crash, no warning. Performing those steps inside the library, where they are + tested, means a caller cannot forget a step they never have to write. + + Args: + positive_texts: Examples of the rare class to detect. + negative_texts: Examples of ordinary, common-class content. + model_name: Sentence transformer to encode with. + neg_to_pos_ratio: Optional negatives-to-positives ratio. None keeps every + negative given. + batch_size: Encoding batch size. + seed: Optional seed for the negative downsampling, so the resulting index + is reproducible. + encoding_additional_kwargs: Extra encoding options, merged over + :data:`DEFAULT_ENCODING_KWARGS`. + model_card: Optional metadata describing where the examples came from. + show_progress_bar: Whether to show the encoder progress bar. + + Returns: + A ready-to-use SentinelLocalIndex, corpus included. + + Raises: + ValueError: If either text list is empty, if a bare string is passed where + a list is expected, or if neg_to_pos_ratio is not positive. + """ + _validate_texts("positive_texts", positive_texts) + _validate_texts("negative_texts", negative_texts) + 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}." + ) + + positive_corpus = list(positive_texts) + negative_corpus = list(negative_texts) + + # Keep both return values: dropping scale_fn silently changes scores for + # models like E5, with nothing to indicate anything went wrong. + sentence_model, scale_fn = get_sentence_transformer_and_scaling_fn(model_name) + + encoding_kwargs = dict(DEFAULT_ENCODING_KWARGS) + encoding_kwargs.update(encoding_additional_kwargs or {}) + + LOG.info( + "Encoding %d positive and %d negative examples with %s", + len(positive_corpus), + len(negative_corpus), + model_name, + ) + positive_embeddings = torch.tensor( + sentence_model.encode( + positive_corpus, + batch_size=batch_size, + show_progress_bar=show_progress_bar, + **encoding_kwargs, + ) + ) + negative_embeddings = torch.tensor( + sentence_model.encode( + negative_corpus, + batch_size=batch_size, + show_progress_bar=show_progress_bar, + **encoding_kwargs, + ) + ) + + if neg_to_pos_ratio is not None: + n_keep = max(1, int(len(positive_corpus) * neg_to_pos_ratio)) + if n_keep < negative_embeddings.shape[0]: + LOG.info( + "Keeping %d negative examples out of %d to reach a %.2f:1 ratio", + n_keep, + negative_embeddings.shape[0], + neg_to_pos_ratio, + ) + generator = ( + torch.Generator().manual_seed(seed) if seed is not None else None + ) + indices = torch.randperm( + negative_embeddings.shape[0], generator=generator + )[:n_keep] + indices = torch.sort(indices).values + negative_embeddings, negative_corpus = _take_rows( + negative_embeddings, negative_corpus, indices + ) + + return cls( + sentence_model=sentence_model, + positive_embeddings=positive_embeddings, + negative_embeddings=negative_embeddings, + scale_fn=scale_fn, + encoding_additional_kwargs=encoding_kwargs, + positive_corpus=positive_corpus, + negative_corpus=negative_corpus, + model_card=model_card, + ) + @classmethod def load( cls, diff --git a/tests/test_from_texts.py b/tests/test_from_texts.py new file mode 100644 index 0000000..2000689 --- /dev/null +++ b/tests/test_from_texts.py @@ -0,0 +1,157 @@ +# Copyright 2025 Roblox Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for SentinelLocalIndex.from_texts().""" + +import tempfile +import pytest +import torch + +from sentinel.sentinel_local_index import SentinelLocalIndex +from sentinel.score_types import RareClassAffinityResult + + +class TestFromTexts: + """Building an index in one call.""" + + POSITIVE = ["unsafe content detected", "harmful behavior observed", "dangerous activity"] + NEGATIVE = [ + "normal behavior detected", + "regular activity observed", + "safe content identified", + "standard procedure followed", + "ordinary events occurred", + "the meeting went well", + ] + + @pytest.mark.integration + def test_builds_a_usable_index(self): + """One call produces an index that scores text correctly.""" + index = SentinelLocalIndex.from_texts( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name="sentence-transformers/all-MiniLM-L6-v2", + ) + + assert index.positive_embeddings.shape[0] == len(self.POSITIVE) + assert index.negative_embeddings.shape[0] == len(self.NEGATIVE) + assert index.sentence_model is not None + + result = index.calculate_rare_class_affinity( + ["harmful unsafe behavior", "normal regular activity"] + ) + assert isinstance(result, RareClassAffinityResult) + + @pytest.mark.integration + def test_corpus_is_always_kept(self): + """The corpus comes along automatically - half the point of the method. + + Forgetting it in the manual recipe costs you explanations, silently. + """ + index = SentinelLocalIndex.from_texts( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name="sentence-transformers/all-MiniLM-L6-v2", + ) + + assert index.positive_corpus == self.POSITIVE + assert index.negative_corpus == self.NEGATIVE + + @pytest.mark.integration + def test_normalization_is_applied_by_default(self): + """Embeddings come out unit-length without the caller asking. + + Omitting normalize_embeddings by hand does not error; it just makes the + similarity maths wrong. Asserting the norms catches a silent regression. + """ + index = SentinelLocalIndex.from_texts( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name="sentence-transformers/all-MiniLM-L6-v2", + ) + + norms = index.positive_embeddings.norm(dim=1) + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) + assert index.encoding_kwargs["normalize_embeddings"] is True + + @pytest.mark.integration + def test_ratio_downsamples_and_keeps_alignment(self): + """The ratio is applied, and surviving negatives keep their own text.""" + index = SentinelLocalIndex.from_texts( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name="sentence-transformers/all-MiniLM-L6-v2", + neg_to_pos_ratio=1.0, + seed=42, + ) + + assert index.negative_embeddings.shape[0] == 3 # 3 positives * 1.0 + assert len(index.negative_corpus) == 3 + assert set(index.negative_corpus) <= set(self.NEGATIVE) + + @pytest.mark.integration + def test_seeded_ratio_is_reproducible(self): + """Same seed, same index.""" + kwargs = dict( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name="sentence-transformers/all-MiniLM-L6-v2", + neg_to_pos_ratio=1.0, + ) + a = SentinelLocalIndex.from_texts(seed=7, **kwargs) + b = SentinelLocalIndex.from_texts(seed=7, **kwargs) + + assert torch.equal(a.negative_embeddings, b.negative_embeddings) + assert a.negative_corpus == b.negative_corpus + + @pytest.mark.integration + def test_round_trip_through_save_and_load(self): + """An index built this way saves and reloads with explanations intact.""" + model_name = "sentence-transformers/all-MiniLM-L6-v2" + index = SentinelLocalIndex.from_texts( + positive_texts=self.POSITIVE, + negative_texts=self.NEGATIVE, + model_name=model_name, + ) + + with tempfile.TemporaryDirectory() as temp_dir: + index.save(path=temp_dir, encoder_model_name_or_path=model_name) + reloaded = SentinelLocalIndex.load( + path=temp_dir, negative_to_positive_ratio=None, seed=1 + ) + + assert reloaded.positive_corpus == self.POSITIVE + assert reloaded.negative_corpus == self.NEGATIVE + + @pytest.mark.parametrize( + "kwargs,message", + [ + ({"positive_texts": "a bare string"}, "positive_texts must be a list"), + ({"negative_texts": "a bare string"}, "negative_texts must be a list"), + ({"positive_texts": []}, "positive_texts must not be empty"), + ({"negative_texts": []}, "negative_texts must not be empty"), + ({"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_input_validation(self, kwargs, message): + """Bad input is rejected up front, before any expensive encoding happens. + + A bare string is iterable, so without this check it would be encoded one + character at a time - confusing, slow, and entirely silent. + """ + call = {"positive_texts": self.POSITIVE, "negative_texts": self.NEGATIVE} + call.update(kwargs) + with pytest.raises(ValueError, match=message): + SentinelLocalIndex.from_texts(**call) From 093a174705f0f35eceee6b0d1b114967a497c172 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:19:34 -0700 Subject: [PATCH 2/2] Drop surplus negatives before encoding, and share the row-selection logic Two review points on from_texts, addressed together because the fix for one decides the shape of the other. from_texts encoded every negative it was given and only then applied the ratio, so anything above the ratio was paid for and thrown away. Encoding is per-text and is the only expensive step here: at a 1:1 ratio against 1,000 positives, a caller passing 100,000 negatives paid to encode 99,000 rows that never reached the index. The choice of which negatives to keep depends only on the positive count and the ratio, both known before encoding, and encoding is per-text, so selecting first yields the same index for a fraction of the cost. Selecting first means selecting texts rather than embedding rows, so the piece worth sharing with subsample() is the choice of positions, not the copy. That is now _choose_indices, used by both, with _select_subset built on top of it for the embedding case. _select_subset also moves to module level, since it never used self and a classmethod cannot reach an instance method. Consolidating matters beyond tidiness: keeping each embedding beside its own text is the one rule in this module that fails silently when broken, so it should live in one place. Adds tests asserting that only the kept negatives ever reach the encoder, and that the no-ratio path still encodes everything. The first fails on the previous encode-then-discard order. _apply_negative_ratio has a third copy of this logic, left alone deliberately: it mutates in place and warns rather than raising, so folding it in would mean re-reviewing merged load behaviour from this PR. Reported in review by vcai4071. Co-authored-by: Cursor --- src/sentinel/sentinel_local_index.py | 152 ++++++++++++++++----------- tests/test_from_texts.py | 63 +++++++++++ 2 files changed, 151 insertions(+), 64 deletions(-) diff --git a/src/sentinel/sentinel_local_index.py b/src/sentinel/sentinel_local_index.py index e293757..e101308 100644 --- a/src/sentinel/sentinel_local_index.py +++ b/src/sentinel/sentinel_local_index.py @@ -129,6 +129,74 @@ def _split_generators( ) +def _choose_indices( + available: int, + n_keep: Optional[int], + generator: Optional[torch.Generator], + label: str, +) -> Optional[torch.Tensor]: + """Pick which ``n_keep`` of ``available`` positions to keep, in their original order. + + Only the choice lives here, not what is done with it, because the two callers + keep different things: :meth:`SentinelLocalIndex.subsample` selects rows of an + existing embedding tensor, while :meth:`SentinelLocalIndex.from_texts` selects + raw texts before paying to encode them. + + Args: + available: How many positions there are to choose from. + n_keep: How many to keep. None, or at least ``available``, keeps everything. + generator: Optional seeded generator for a reproducible choice. + label: "positive" or "negative", used in log messages. + + Returns: + Sorted positions to keep, or None when everything is kept - which lets + callers skip the copy entirely rather than rebuild an identical list. + """ + 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 only %d are available - keeping all of them.", + n_keep, + label, + available, + ) + return 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 indices + + +def _select_subset( + 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 an 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. + """ + indices = _choose_indices(embeddings.shape[0], n_keep, generator, label) + if indices is None: + # Copy the corpus list so callers cannot mutate the original through the copy. + return embeddings, (list(corpus) if corpus is not None else None) + return _take_rows(embeddings, corpus, indices) + + def _take_rows( embeddings: torch.Tensor, corpus: Optional[List[str]], @@ -322,7 +390,8 @@ def from_texts( negative_texts: Examples of ordinary, common-class content. model_name: Sentence transformer to encode with. neg_to_pos_ratio: Optional negatives-to-positives ratio. None keeps every - negative given. + negative given. Surplus negatives are dropped before encoding, so + passing far more than the ratio needs costs little. batch_size: Encoding batch size. seed: Optional seed for the negative downsampling, so the resulting index is reproducible. @@ -355,6 +424,22 @@ def from_texts( encoding_kwargs = dict(DEFAULT_ENCODING_KWARGS) encoding_kwargs.update(encoding_additional_kwargs or {}) + # Drop the surplus negatives before encoding, not after. Encoding is the only + # expensive step here, and it is per-text, so encoding a sentence and then + # discarding it is pure waste: at a 1:1 ratio against 1,000 positives, a + # caller passing 100,000 negatives would have paid to encode 99,000 rows + # that never reach the index. + if neg_to_pos_ratio is not None: + n_keep = max(1, int(len(positive_corpus) * neg_to_pos_ratio)) + generator = ( + torch.Generator().manual_seed(seed) if seed is not None else None + ) + indices = _choose_indices( + len(negative_corpus), n_keep, generator, "negative" + ) + if indices is not None: + negative_corpus = [negative_corpus[i] for i in indices.tolist()] + LOG.info( "Encoding %d positive and %d negative examples with %s", len(positive_corpus), @@ -378,26 +463,6 @@ def from_texts( ) ) - if neg_to_pos_ratio is not None: - n_keep = max(1, int(len(positive_corpus) * neg_to_pos_ratio)) - if n_keep < negative_embeddings.shape[0]: - LOG.info( - "Keeping %d negative examples out of %d to reach a %.2f:1 ratio", - n_keep, - negative_embeddings.shape[0], - neg_to_pos_ratio, - ) - generator = ( - torch.Generator().manual_seed(seed) if seed is not None else None - ) - indices = torch.randperm( - negative_embeddings.shape[0], generator=generator - )[:n_keep] - indices = torch.sort(indices).values - negative_embeddings, negative_corpus = _take_rows( - negative_embeddings, negative_corpus, indices - ) - return cls( sentence_model=sentence_model, positive_embeddings=positive_embeddings, @@ -567,47 +632,6 @@ 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, @@ -661,7 +685,7 @@ def subsample( # 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( + positive_embeddings, positive_corpus = _select_subset( self.positive_embeddings, self.positive_corpus, n_positive, @@ -681,7 +705,7 @@ def subsample( ) n_negative = 1 - negative_embeddings, negative_corpus = self._select_subset( + negative_embeddings, negative_corpus = _select_subset( self.negative_embeddings, self.negative_corpus, n_negative, diff --git a/tests/test_from_texts.py b/tests/test_from_texts.py index 2000689..0353920 100644 --- a/tests/test_from_texts.py +++ b/tests/test_from_texts.py @@ -15,6 +15,7 @@ """Tests for SentinelLocalIndex.from_texts().""" import tempfile +import numpy as np import pytest import torch @@ -22,6 +23,21 @@ from sentinel.score_types import RareClassAffinityResult +class _SpyModel: + """Stands in for a sentence transformer, recording what it was asked to encode. + + Encoding is the only expensive step in from_texts, so what reaches this stub is + exactly what a real run would pay for. + """ + + def __init__(self): + self.encoded = [] + + def encode(self, texts, **kwargs): + self.encoded.append(list(texts)) + return np.zeros((len(texts), 4), dtype=np.float32) + + class TestFromTexts: """Building an index in one call.""" @@ -134,6 +150,53 @@ def test_round_trip_through_save_and_load(self): assert reloaded.positive_corpus == self.POSITIVE assert reloaded.negative_corpus == self.NEGATIVE + def _spy_model(self, monkeypatch): + """Swap the real encoder for a spy, and hand the spy back.""" + spy = _SpyModel() + monkeypatch.setattr( + "sentinel.sentinel_local_index.get_sentence_transformer_and_scaling_fn", + lambda *args, **kwargs: (spy, None), + ) + return spy + + def test_surplus_negatives_are_never_encoded(self, monkeypatch): + """The ratio must be applied before encoding, not after. + + Encoding is per-text and is the only expensive step, so encoding a negative + and then discarding it is pure waste. At a 1:1 ratio against 3 positives, + only 3 of the 50 negatives belong in the index, so only 3 should ever reach + the encoder. + """ + spy = self._spy_model(monkeypatch) + positives = ["p0", "p1", "p2"] + negatives = [f"n{i}" for i in range(50)] + + index = SentinelLocalIndex.from_texts( + positive_texts=positives, + negative_texts=negatives, + neg_to_pos_ratio=1.0, + seed=42, + ) + + assert [len(call) for call in spy.encoded] == [3, 3] + assert spy.encoded[0] == positives + # Everything encoded ended up in the index, and vice versa: nothing was + # paid for and thrown away. + assert spy.encoded[1] == index.negative_corpus + + def test_without_a_ratio_every_negative_is_encoded(self, monkeypatch): + """Reordering the downsample must not change the no-ratio path.""" + spy = self._spy_model(monkeypatch) + negatives = [f"n{i}" for i in range(20)] + + index = SentinelLocalIndex.from_texts( + positive_texts=["p0", "p1"], + negative_texts=negatives, + ) + + assert spy.encoded[1] == negatives + assert index.negative_corpus == negatives + @pytest.mark.parametrize( "kwargs,message", [