|
| 1 | +# Copyright 2025 Roblox Corporation |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Tests for SentinelLocalIndex.from_texts().""" |
| 16 | + |
| 17 | +import tempfile |
| 18 | +import pytest |
| 19 | +import torch |
| 20 | + |
| 21 | +from sentinel.sentinel_local_index import SentinelLocalIndex |
| 22 | +from sentinel.score_types import RareClassAffinityResult |
| 23 | + |
| 24 | + |
| 25 | +class TestFromTexts: |
| 26 | + """Building an index in one call.""" |
| 27 | + |
| 28 | + POSITIVE = ["unsafe content detected", "harmful behavior observed", "dangerous activity"] |
| 29 | + NEGATIVE = [ |
| 30 | + "normal behavior detected", |
| 31 | + "regular activity observed", |
| 32 | + "safe content identified", |
| 33 | + "standard procedure followed", |
| 34 | + "ordinary events occurred", |
| 35 | + "the meeting went well", |
| 36 | + ] |
| 37 | + |
| 38 | + @pytest.mark.integration |
| 39 | + def test_builds_a_usable_index(self): |
| 40 | + """One call produces an index that scores text correctly.""" |
| 41 | + index = SentinelLocalIndex.from_texts( |
| 42 | + positive_texts=self.POSITIVE, |
| 43 | + negative_texts=self.NEGATIVE, |
| 44 | + model_name="sentence-transformers/all-MiniLM-L6-v2", |
| 45 | + ) |
| 46 | + |
| 47 | + assert index.positive_embeddings.shape[0] == len(self.POSITIVE) |
| 48 | + assert index.negative_embeddings.shape[0] == len(self.NEGATIVE) |
| 49 | + assert index.sentence_model is not None |
| 50 | + |
| 51 | + result = index.calculate_rare_class_affinity( |
| 52 | + ["harmful unsafe behavior", "normal regular activity"] |
| 53 | + ) |
| 54 | + assert isinstance(result, RareClassAffinityResult) |
| 55 | + |
| 56 | + @pytest.mark.integration |
| 57 | + def test_corpus_is_always_kept(self): |
| 58 | + """The corpus comes along automatically - half the point of the method. |
| 59 | +
|
| 60 | + Forgetting it in the manual recipe costs you explanations, silently. |
| 61 | + """ |
| 62 | + index = SentinelLocalIndex.from_texts( |
| 63 | + positive_texts=self.POSITIVE, |
| 64 | + negative_texts=self.NEGATIVE, |
| 65 | + model_name="sentence-transformers/all-MiniLM-L6-v2", |
| 66 | + ) |
| 67 | + |
| 68 | + assert index.positive_corpus == self.POSITIVE |
| 69 | + assert index.negative_corpus == self.NEGATIVE |
| 70 | + |
| 71 | + @pytest.mark.integration |
| 72 | + def test_normalization_is_applied_by_default(self): |
| 73 | + """Embeddings come out unit-length without the caller asking. |
| 74 | +
|
| 75 | + Omitting normalize_embeddings by hand does not error; it just makes the |
| 76 | + similarity maths wrong. Asserting the norms catches a silent regression. |
| 77 | + """ |
| 78 | + index = SentinelLocalIndex.from_texts( |
| 79 | + positive_texts=self.POSITIVE, |
| 80 | + negative_texts=self.NEGATIVE, |
| 81 | + model_name="sentence-transformers/all-MiniLM-L6-v2", |
| 82 | + ) |
| 83 | + |
| 84 | + norms = index.positive_embeddings.norm(dim=1) |
| 85 | + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) |
| 86 | + assert index.encoding_kwargs["normalize_embeddings"] is True |
| 87 | + |
| 88 | + @pytest.mark.integration |
| 89 | + def test_ratio_downsamples_and_keeps_alignment(self): |
| 90 | + """The ratio is applied, and surviving negatives keep their own text.""" |
| 91 | + index = SentinelLocalIndex.from_texts( |
| 92 | + positive_texts=self.POSITIVE, |
| 93 | + negative_texts=self.NEGATIVE, |
| 94 | + model_name="sentence-transformers/all-MiniLM-L6-v2", |
| 95 | + neg_to_pos_ratio=1.0, |
| 96 | + seed=42, |
| 97 | + ) |
| 98 | + |
| 99 | + assert index.negative_embeddings.shape[0] == 3 # 3 positives * 1.0 |
| 100 | + assert len(index.negative_corpus) == 3 |
| 101 | + assert set(index.negative_corpus) <= set(self.NEGATIVE) |
| 102 | + |
| 103 | + @pytest.mark.integration |
| 104 | + def test_seeded_ratio_is_reproducible(self): |
| 105 | + """Same seed, same index.""" |
| 106 | + kwargs = dict( |
| 107 | + positive_texts=self.POSITIVE, |
| 108 | + negative_texts=self.NEGATIVE, |
| 109 | + model_name="sentence-transformers/all-MiniLM-L6-v2", |
| 110 | + neg_to_pos_ratio=1.0, |
| 111 | + ) |
| 112 | + a = SentinelLocalIndex.from_texts(seed=7, **kwargs) |
| 113 | + b = SentinelLocalIndex.from_texts(seed=7, **kwargs) |
| 114 | + |
| 115 | + assert torch.equal(a.negative_embeddings, b.negative_embeddings) |
| 116 | + assert a.negative_corpus == b.negative_corpus |
| 117 | + |
| 118 | + @pytest.mark.integration |
| 119 | + def test_round_trip_through_save_and_load(self): |
| 120 | + """An index built this way saves and reloads with explanations intact.""" |
| 121 | + model_name = "sentence-transformers/all-MiniLM-L6-v2" |
| 122 | + index = SentinelLocalIndex.from_texts( |
| 123 | + positive_texts=self.POSITIVE, |
| 124 | + negative_texts=self.NEGATIVE, |
| 125 | + model_name=model_name, |
| 126 | + ) |
| 127 | + |
| 128 | + with tempfile.TemporaryDirectory() as temp_dir: |
| 129 | + index.save(path=temp_dir, encoder_model_name_or_path=model_name) |
| 130 | + reloaded = SentinelLocalIndex.load( |
| 131 | + path=temp_dir, negative_to_positive_ratio=None, seed=1 |
| 132 | + ) |
| 133 | + |
| 134 | + assert reloaded.positive_corpus == self.POSITIVE |
| 135 | + assert reloaded.negative_corpus == self.NEGATIVE |
| 136 | + |
| 137 | + @pytest.mark.parametrize( |
| 138 | + "kwargs,message", |
| 139 | + [ |
| 140 | + ({"positive_texts": "a bare string"}, "positive_texts must be a list"), |
| 141 | + ({"negative_texts": "a bare string"}, "negative_texts must be a list"), |
| 142 | + ({"positive_texts": []}, "positive_texts must not be empty"), |
| 143 | + ({"negative_texts": []}, "negative_texts must not be empty"), |
| 144 | + ({"neg_to_pos_ratio": 0}, "neg_to_pos_ratio must be positive"), |
| 145 | + ({"neg_to_pos_ratio": -1.0}, "neg_to_pos_ratio must be positive"), |
| 146 | + ], |
| 147 | + ) |
| 148 | + def test_input_validation(self, kwargs, message): |
| 149 | + """Bad input is rejected up front, before any expensive encoding happens. |
| 150 | +
|
| 151 | + A bare string is iterable, so without this check it would be encoded one |
| 152 | + character at a time - confusing, slow, and entirely silent. |
| 153 | + """ |
| 154 | + call = {"positive_texts": self.POSITIVE, "negative_texts": self.NEGATIVE} |
| 155 | + call.update(kwargs) |
| 156 | + with pytest.raises(ValueError, match=message): |
| 157 | + SentinelLocalIndex.from_texts(**call) |
0 commit comments