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..e101308 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] @@ -102,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]], @@ -214,9 +309,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 +364,116 @@ 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. 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. + 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 {}) + + # 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), + 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, + ) + ) + + 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, @@ -429,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, @@ -523,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, @@ -543,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 new file mode 100644 index 0000000..0353920 --- /dev/null +++ b/tests/test_from_texts.py @@ -0,0 +1,220 @@ +# 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 numpy as np +import pytest +import torch + +from sentinel.sentinel_local_index import SentinelLocalIndex +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.""" + + 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 + + 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", + [ + ({"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)