Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 43 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
144 changes: 141 additions & 3 deletions src/sentinel/sentinel_local_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
157 changes: 157 additions & 0 deletions tests/test_from_texts.py
Original file line number Diff line number Diff line change
@@ -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)
Loading