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 6b73bd1..53acdd6 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] @@ -182,9 +209,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 @@ -239,6 +264,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)