diff --git a/README.md b/README.md index 8cfc030..0c55e63 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,11 @@ from sentinel.sentinel_local_index import SentinelLocalIndex # Load a previously saved index from a local path index = SentinelLocalIndex.load(path="path/to/local/index") +# Loading downsamples the negatives to the requested ratio, and that choice is +# random. Pass a seed to make it reproducible, so the same saved index always +# behaves identically - worth doing whenever you are tuning or debugging. +index = SentinelLocalIndex.load(path="path/to/local/index", seed=42) + # Or load from S3 index = SentinelLocalIndex.load( path="s3://my-bucket/path/to/index", @@ -338,6 +343,19 @@ 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: + +| 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` 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. + ## Examples To run the notebook examples ```bash diff --git a/src/sentinel/io/__init__.py b/src/sentinel/io/__init__.py index 27a68ae..219ac64 100644 --- a/src/sentinel/io/__init__.py +++ b/src/sentinel/io/__init__.py @@ -15,6 +15,6 @@ """IO utilities for Sentinel.""" from sentinel.io.saved_index_config import SavedIndexConfig -from sentinel.io.index_io import save_index, load_index +from sentinel.io.index_io import save_index, load_index, load_corpus -__all__ = ["SavedIndexConfig", "save_index", "load_index"] +__all__ = ["SavedIndexConfig", "save_index", "load_index", "load_corpus"] diff --git a/src/sentinel/io/index_io.py b/src/sentinel/io/index_io.py index 12952c9..a587af0 100644 --- a/src/sentinel/io/index_io.py +++ b/src/sentinel/io/index_io.py @@ -23,7 +23,7 @@ import logging import tempfile import os -from typing import Dict, Any, Optional, Tuple +from typing import Dict, Any, List, Optional, Tuple import torch import safetensors @@ -37,9 +37,14 @@ # Constants for file names CONFIG_FILE_NAME = "sentinel_local_index_config.json" EMBEDDINGS_FILE_NAME = "embeddings.safetensors" +# Always written by save_index, holding nulls when no corpus is supplied. Absent +# only in indices saved before corpus support existed. +CORPUS_FILE_NAME = "corpus.json" # Storage keys - kept as positive/negative for backward compatibility POSITIVE_EMBEDDINGS_KEY = "positive_embeddings" # Corresponds to rare class examples NEGATIVE_EMBEDDINGS_KEY = "negative_embeddings" # Corresponds to common class examples +POSITIVE_CORPUS_KEY = "positive_corpus" +NEGATIVE_CORPUS_KEY = "negative_corpus" def create_s3_transport_params( @@ -87,12 +92,38 @@ def _join_path(base_path: str, filename: str) -> str: return os.path.join(base_path, filename) +def _check_corpus_length( + name: str, corpus: Optional[List[str]], embeddings: Optional[torch.Tensor] +) -> None: + """Raise if a corpus does not line up row-for-row with its embeddings. + + Args: + name: Which corpus is being checked, used in the error message. + corpus: The corpus texts, or None when there is nothing to check. + embeddings: The embeddings the corpus is supposed to describe. + + Raises: + ValueError: If the corpus and embeddings have different lengths. + """ + if corpus is None or embeddings is None: + return + if len(corpus) != embeddings.shape[0]: + raise ValueError( + f"{name}_corpus has {len(corpus)} entries but {name}_embeddings has " + f"{embeddings.shape[0]} rows. Saving a misaligned corpus would make " + f"explanations name the wrong text, so refusing to write it." + ) + + def save_index( path: str, config: SavedIndexConfig, positive_embeddings: torch.Tensor, # Represents rare class examples negative_embeddings: torch.Tensor, # Represents common class examples transport_params: Optional[Dict[str, Any]] = None, + *, + positive_corpus: Optional[List[str]] = None, + negative_corpus: Optional[List[str]] = None, ) -> None: """ Save a SentinelLocalIndex to a path. @@ -103,7 +134,23 @@ def save_index( positive_embeddings: Tensor of positive (rare class) example embeddings. negative_embeddings: Tensor of negative (common class) example embeddings. transport_params: Optional transport parameters for smart_open. + positive_corpus: Optional texts behind the positive embeddings. Written to + an extra ``corpus.json`` so explanations survive a reload. + negative_corpus: Optional texts behind the negative embeddings. + + Note: + ``corpus.json`` is always written, storing nulls when a corpus is not + supplied, so that it never describes rows from an earlier save to the same + path. Saving without a corpus therefore clears any corpus already there. + + Raises: + ValueError: If a corpus is supplied whose length does not match its embeddings. """ + # Fail before writing anything: a misaligned corpus baked into the artifact is + # worse than no corpus, and the caller is right here and can fix it. + _check_corpus_length("positive", positive_corpus, positive_embeddings) + _check_corpus_length("negative", negative_corpus, negative_embeddings) + # Ensure directory exists for local paths if not path.startswith("s3://") and not os.path.exists(path): os.makedirs(path, exist_ok=True) @@ -146,6 +193,25 @@ def save_index( ) as f_out: f_out.write(f_in.read()) + # Written unconditionally, nulls included. Skipping the write when no corpus is + # supplied would leave an earlier save's corpus.json next to the embeddings just + # written, and a stale corpus that happens to have the same row count passes + # every alignment check we have, so explanations would name the wrong text with + # no warning. Writing nulls keeps the file describing the embeddings beside it. + # + # The corpus is small text rather than a tensor blob, so smart_open handles both + # local and S3 destinations directly, the same way the config file does. + corpus_path = _join_path(path, CORPUS_FILE_NAME) + LOG.info("Saving corpus to %s", corpus_path) + with smart_open.open(corpus_path, "w", transport_params=transport_params) as f: + json.dump( + { + POSITIVE_CORPUS_KEY: positive_corpus, + NEGATIVE_CORPUS_KEY: negative_corpus, + }, + f, + ) + LOG.info("Successfully saved index to %s", path) @@ -209,6 +275,49 @@ def _load_embeddings_from_file(file_path: str) -> Tuple[torch.Tensor, torch.Tens return positive_embeddings, negative_embeddings +def load_corpus( + path: str, transport_params: Optional[Dict[str, Any]] = None +) -> Tuple[Optional[List[str]], Optional[List[str]]]: + """ + Load the corpus texts of an index, if it has any. + + This is a separate function rather than extra return values on + :func:`load_index` because ``load_index`` is public API that callers unpack as a + 3-tuple; widening it would break them. The corpus is an optional, separate file, + so an optional, separate loader mirrors the format honestly. + + Args: + path: Path where the index is stored. + transport_params: Optional transport parameters for smart_open. + + Returns: + Tuple of (positive_corpus, negative_corpus). Either element may be None, + whether because the file stores an explicit null or because the file is + absent entirely, as it is for indices saved before corpus support. + + Raises: + json.JSONDecodeError: If the corpus file exists but is not valid JSON. A + corrupt artifact is a real problem and is not silently swallowed. + """ + corpus_path = _join_path(path, CORPUS_FILE_NAME) + + try: + with smart_open.open(corpus_path, "r", transport_params=transport_params) as f: + payload = json.load(f) + except (FileNotFoundError, OSError, KeyError) as e: + # Every index saved before corpus support lacks this file. That is the normal + # case, not an error. smart_open surfaces a missing S3 key as OSError/KeyError. + LOG.info( + "No corpus file at %s (%s) - explanations will fall back to row numbers.", + corpus_path, + type(e).__name__, + ) + return None, None + + LOG.info("Loaded corpus from %s", corpus_path) + return payload.get(POSITIVE_CORPUS_KEY), payload.get(NEGATIVE_CORPUS_KEY) + + def load_index( path: str, transport_params: Optional[Dict[str, Any]] = None ) -> Tuple[SavedIndexConfig, torch.Tensor, torch.Tensor]: diff --git a/src/sentinel/sentinel_local_index.py b/src/sentinel/sentinel_local_index.py index 3a67000..6b73bd1 100644 --- a/src/sentinel/sentinel_local_index.py +++ b/src/sentinel/sentinel_local_index.py @@ -19,7 +19,7 @@ """ import logging -from typing import Optional, List, Mapping, Any, Callable +from typing import Optional, List, Mapping, Any, Callable, Tuple import numpy as np import torch @@ -28,13 +28,81 @@ from sentinel.score_formulae import calculate_contrastive_score, skewness, contrastive_components from sentinel.io.saved_index_config import SavedIndexConfig -from sentinel.io.index_io import save_index, load_index, create_s3_transport_params +from sentinel.io.index_io import ( + save_index, + load_index, + load_corpus, + create_s3_transport_params, +) from sentinel.embeddings.sbert import get_sentence_transformer_and_scaling_fn from sentinel.score_types import RareClassAffinityResult LOG = logging.getLogger(__name__) +def _corpus_if_aligned( + name: str, corpus: Optional[List[str]], embeddings: Optional[torch.Tensor] +) -> Optional[List[str]]: + """Return the corpus only if it lines up with its embeddings, else None. + + Degrading to row numbers is recoverable; confidently reporting the wrong + sentence as the reason for a score is not, so a mismatch is discarded. + + Args: + name: Which corpus is being checked, used in the warning. + corpus: The corpus texts, or None. + embeddings: The embeddings the corpus is supposed to describe. + + Returns: + The corpus unchanged, or None if it is absent or misaligned. + """ + if corpus is None or embeddings is None: + return None + if len(corpus) != embeddings.shape[0]: + LOG.warning( + "Discarding %s corpus: %d texts for %d embedding rows. Explanations will " + "show row numbers instead of naming the wrong text.", + name, + len(corpus), + embeddings.shape[0], + ) + return None + return corpus + + +def _take_rows( + embeddings: torch.Tensor, + corpus: Optional[List[str]], + indices: torch.Tensor, +) -> Tuple[torch.Tensor, Optional[List[str]]]: + """Select rows from embeddings and the matching corpus texts, keeping them aligned. + + The embeddings and the corpus are two parallel lists: row *i* of one describes + entry *i* of the other. Any operation that drops rows has to drop the same + positions from both, or explanations will confidently name unrelated text. + + Args: + embeddings: Tensor to select rows from. + corpus: Texts describing those rows, or None if the index has no corpus. + indices: Row positions to keep. + + Returns: + Tuple of (selected embeddings, selected corpus or None). + """ + selected = embeddings[indices] + if corpus is None: + return selected, None + if len(corpus) != embeddings.shape[0]: + LOG.warning( + "Corpus length %d does not match embedding rows %d - dropping corpus " + "rather than risk misaligned explanations.", + len(corpus), + embeddings.shape[0], + ) + return selected, None + return selected, [corpus[i] for i in indices.tolist()] + + class SentinelLocalIndex: """Calculate scores for detecting extremely rare classes of text using contrastive learning. @@ -164,6 +232,8 @@ def save( positive_embeddings=self.positive_embeddings, negative_embeddings=self.negative_embeddings, transport_params=transport_params, + positive_corpus=self.positive_corpus, + negative_corpus=self.negative_corpus, ) # Return the config for informational purposes @@ -177,6 +247,7 @@ def load( aws_secret_access_key: Optional[str] = None, negative_to_positive_ratio: Optional[float] = 5.0, cache_model: bool = False, + seed: Optional[int] = None, ) -> "SentinelLocalIndex": """ Load the index from a path and returns a new SentinelLocalIndex instance. @@ -190,6 +261,9 @@ def load( If 5.0 (default), uses a 5:1 negative to positive ratio for optimal performance. If specified, downsamples negative examples to achieve the desired ratio. cache_model: Whether to use model caching for faster subsequent loads. Default True. + seed: Optional seed for the negative downsampling. Without it the surviving + subset differs on every load, so the same saved index scores slightly + differently each time. Pass a seed to make loading reproducible. Returns: A new SentinelLocalIndex instance with the loaded model and embeddings. @@ -204,6 +278,17 @@ def load( path=path, transport_params=transport_params ) + # Optional; absent for indices saved before corpus support. + positive_corpus, negative_corpus = load_corpus( + path=path, transport_params=transport_params + ) + positive_corpus = _corpus_if_aligned( + "positive", positive_corpus, positive_embeddings + ) + negative_corpus = _corpus_if_aligned( + "negative", negative_corpus, negative_embeddings + ) + # Create the sentence model and get the scaling function model_name = config.encoder_model_name_or_path @@ -219,15 +304,19 @@ def load( positive_embeddings=positive_embeddings, negative_embeddings=negative_embeddings, encoding_additional_kwargs=config.encoding_kwargs, + positive_corpus=positive_corpus, + negative_corpus=negative_corpus, model_card=config.model_card, ) # Apply negative ratio if needed - instance._apply_negative_ratio(negative_to_positive_ratio) + instance._apply_negative_ratio(negative_to_positive_ratio, seed=seed) return instance - def _apply_negative_ratio(self, negative_to_positive_ratio: Optional[float]): + def _apply_negative_ratio( + self, negative_to_positive_ratio: Optional[float], seed: Optional[int] = None + ): """ Apply the negative_to_positive_ratio to reduce the number of negative (common class) examples. @@ -235,6 +324,9 @@ def _apply_negative_ratio(self, negative_to_positive_ratio: Optional[float]): negative_to_positive_ratio: The ratio of negative samples to keep relative to positive samples. If None, preserves the original ratio from the saved index. If 5.0 (default), uses optimized 5:1 ratio for best performance. + seed: Optional seed making the choice of surviving negatives reproducible. + Uses a private torch.Generator rather than a global seed, so nothing + else in the caller's program has its randomness reset. """ # Handle null/invalid inputs - preserve original ratio if any issues occur if negative_to_positive_ratio is None: @@ -283,10 +375,18 @@ def _apply_negative_ratio(self, negative_to_positive_ratio: Optional[float]): ) # Randomly select a subset of the negative examples with error handling try: - indices = torch.randperm(self.negative_embeddings.shape[0])[ - :num_negative_to_keep - ] - self.negative_embeddings = self.negative_embeddings[indices] + generator = ( + torch.Generator().manual_seed(seed) if seed is not None else None + ) + indices = torch.randperm( + self.negative_embeddings.shape[0], generator=generator + )[:num_negative_to_keep] + # Sorting does not change which rows survive, only their order. Keeping + # the original relative order makes the result far easier to diff and debug. + indices = torch.sort(indices).values + self.negative_embeddings, self.negative_corpus = _take_rows( + self.negative_embeddings, self.negative_corpus, indices + ) except (RuntimeError, IndexError, TypeError) as e: LOG.error("Error during negative embedding downsampling: %s. Preserving original embeddings.", str(e)) return diff --git a/tests/test_io.py b/tests/test_io.py index 13d79b6..f6fd7e5 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -24,9 +24,11 @@ from sentinel.io.index_io import ( save_index, load_index, + load_corpus, create_s3_transport_params, CONFIG_FILE_NAME, EMBEDDINGS_FILE_NAME, + CORPUS_FILE_NAME, ) @@ -116,6 +118,156 @@ def test_index_io_local(positive_shape, negative_shape): assert torch.allclose(loaded_negative, negative_embeddings) +def _config(): + """Build a throwaway config for corpus tests.""" + return SavedIndexConfig( + encoder_model_name_or_path="all-MiniLM-L6-v2", + encoding_kwargs={"normalize_embeddings": True}, + ) + + +def test_corpus_round_trip(): + """A saved corpus comes back identical, so explanations survive a reload.""" + positive_corpus = [f"pos {i}" for i in range(6)] + negative_corpus = [f"neg {i}" for i in range(9)] + + with tempfile.TemporaryDirectory() as temp_dir: + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(6, 8), + negative_embeddings=torch.rand(9, 8), + positive_corpus=positive_corpus, + negative_corpus=negative_corpus, + ) + + assert os.path.exists(os.path.join(temp_dir, CORPUS_FILE_NAME)) + + loaded_positive, loaded_negative = load_corpus(path=temp_dir) + assert loaded_positive == positive_corpus + assert loaded_negative == negative_corpus + + +def test_index_without_corpus_still_loads(): + """Saving with no corpus writes explicit nulls, and loading reports no corpus. + + The file is written even with nothing to put in it so that it always describes + the embeddings beside it; see test_saving_without_corpus_clears_stale_corpus. + """ + with tempfile.TemporaryDirectory() as temp_dir: + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(4, 8), + negative_embeddings=torch.rand(4, 8), + ) + + corpus_file = os.path.join(temp_dir, CORPUS_FILE_NAME) + assert os.path.exists(corpus_file) + with open(corpus_file) as f: + assert json.load(f) == {"positive_corpus": None, "negative_corpus": None} + + # Both the existing loader and the new one behave. + config, positive, negative = load_index(path=temp_dir) + assert positive.shape == (4, 8) + assert load_corpus(path=temp_dir) == (None, None) + + +def test_saving_without_corpus_clears_stale_corpus(): + """Re-saving without a corpus must not leave the previous corpus behind. + + The row counts here are deliberately identical, which is the case that slips + past every alignment check: a corpus of 3 texts lines up with 3 new embedding + rows, so nothing downstream can tell it describes the wrong 3 rows. + """ + with tempfile.TemporaryDirectory() as temp_dir: + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(3, 8), + negative_embeddings=torch.rand(3, 8), + positive_corpus=["first", "second", "third"], + negative_corpus=["one", "two", "three"], + ) + assert load_corpus(path=temp_dir) == ( + ["first", "second", "third"], + ["one", "two", "three"], + ) + + # Same shape, different content, and this time no corpus to go with it. + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(3, 8), + negative_embeddings=torch.rand(3, 8), + ) + + assert load_corpus(path=temp_dir) == (None, None) + + +def test_deleting_corpus_file_degrades_gracefully(): + """An index with no corpus.json at all leaves the rest usable. + + This is the backward-compatibility guarantee: indices saved before corpus + support lack the file entirely, and a missing file has to mean "no corpus" + rather than an error. + """ + with tempfile.TemporaryDirectory() as temp_dir: + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(3, 8), + negative_embeddings=torch.rand(3, 8), + positive_corpus=["a", "b", "c"], + negative_corpus=["x", "y", "z"], + ) + os.remove(os.path.join(temp_dir, CORPUS_FILE_NAME)) + + assert load_corpus(path=temp_dir) == (None, None) + # load_index is untouched by any of this. + config, positive, negative = load_index(path=temp_dir) + assert positive.shape == (3, 8) + + +@pytest.mark.parametrize("which", ["positive", "negative"]) +def test_save_refuses_misaligned_corpus(which): + """A corpus that does not match its embeddings is refused at save time. + + Failing here is deliberate: writing it would bake the misalignment into the + artifact, where it would later surface as confidently wrong explanations. + """ + kwargs = { + "positive_embeddings": torch.rand(5, 8), + "negative_embeddings": torch.rand(5, 8), + f"{which}_corpus": ["only", "three", "texts"], + } + + with tempfile.TemporaryDirectory() as temp_dir: + with pytest.raises(ValueError, match=f"{which}_corpus has 3 entries"): + save_index(path=temp_dir, config=_config(), **kwargs) + + # Nothing partial was left behind. + assert not os.path.exists(os.path.join(temp_dir, CORPUS_FILE_NAME)) + + +def test_malformed_corpus_json_propagates(): + """Corrupt JSON is a real problem and is not silently swallowed.""" + with tempfile.TemporaryDirectory() as temp_dir: + save_index( + path=temp_dir, + config=_config(), + positive_embeddings=torch.rand(2, 8), + negative_embeddings=torch.rand(2, 8), + positive_corpus=["a", "b"], + negative_corpus=["c", "d"], + ) + with open(os.path.join(temp_dir, CORPUS_FILE_NAME), "w") as f: + f.write("{not valid json") + + with pytest.raises(json.JSONDecodeError): + load_corpus(path=temp_dir) + + def test_create_s3_transport_params(): """Test creating S3 transport parameters.""" # Test with credentials diff --git a/tests/test_sentinel_local_index.py b/tests/test_sentinel_local_index.py index 57b9009..ebcde29 100644 --- a/tests/test_sentinel_local_index.py +++ b/tests/test_sentinel_local_index.py @@ -256,6 +256,35 @@ def test_end_to_end_workflow(): # Negative examples should be zero assert negative_score == 0, "Negative example should score zero" + # 7. The corpus survives the round trip, so explanations can name real text. + assert new_index.positive_corpus == positive_texts + # A 1:1 ratio against 4 positives keeps 4 of the 5 negatives; whichever + # survive must still be the real texts, not row numbers. + assert new_index.negative_corpus is not None + assert len(new_index.negative_corpus) == 4 + assert set(new_index.negative_corpus) <= set(negative_texts) + + # 8. Explanations name sentences rather than degrading to row indices, which + # is the user-visible symptom of the corpus being dropped on save. + neighbors = [ + record["neighbor"] + for explanation in result.explanations.values() + for record in (explanation["neighbors"] or []) + ] + assert neighbors, "expected at least one neighbor record" + assert all(isinstance(n, str) for n in neighbors) + assert all(n in set(positive_texts) | set(negative_texts) for n in neighbors) + + # 9. Loading with a seed is reproducible, so a saved index is a fixed artifact. + seeded_a = SentinelLocalIndex.load( + path=temp_dir, negative_to_positive_ratio=1.0, seed=42 + ) + seeded_b = SentinelLocalIndex.load( + path=temp_dir, negative_to_positive_ratio=1.0, seed=42 + ) + assert torch.equal(seeded_a.negative_embeddings, seeded_b.negative_embeddings) + assert seeded_a.negative_corpus == seeded_b.negative_corpus + class TestSentinelLocalIndexEdgeCases: """Test edge cases and error handling in SentinelLocalIndex.""" @@ -503,3 +532,108 @@ def test_empty_observation_scores_line_503(self, simple_index): assert isinstance(result, RareClassAffinityResult) assert result.rare_class_affinity_score == 0.0 # Should be 0.0 when no scores pass + + +def _labelled_index(n_positive=4, n_negative=40, dim=8): + """Build an index where embedding row i provably belongs to text "neg i". + + Row i is filled with the value i, so after any reshuffle you can read an + embedding and know which text it must be paired with. That makes misalignment + detectable rather than merely suspected. + """ + negative_embeddings = torch.arange(n_negative, dtype=torch.float32).repeat(dim, 1).T + return SentinelLocalIndex( + sentence_model=None, + positive_embeddings=torch.ones(n_positive, dim), + negative_embeddings=negative_embeddings, + positive_corpus=[f"pos {i}" for i in range(n_positive)], + negative_corpus=[f"neg {i}" for i in range(n_negative)], + ) + + +class TestNegativeDownsampling: + """Seeding and corpus alignment when negatives are downsampled.""" + + def test_same_seed_gives_identical_subset(self): + """A seed makes the surviving negatives reproducible.""" + first = _labelled_index(n_negative=500) + second = _labelled_index(n_negative=500) + + first._apply_negative_ratio(50.0, seed=42) + second._apply_negative_ratio(50.0, seed=42) + + assert torch.equal(first.negative_embeddings, second.negative_embeddings) + assert first.negative_corpus == second.negative_corpus + + def test_different_seeds_give_different_subsets(self): + """Different seeds pick different negatives, confirming the seed is used.""" + first = _labelled_index(n_negative=500) + second = _labelled_index(n_negative=500) + + first._apply_negative_ratio(50.0, seed=1) + second._apply_negative_ratio(50.0, seed=2) + + assert not torch.equal(first.negative_embeddings, second.negative_embeddings) + + def test_unseeded_downsampling_varies(self): + """Without a seed the subset differs per call - the behaviour being fixed. + + 500 rows choosing 200 makes an accidental collision effectively impossible, + so this is not a flaky assertion. + """ + first = _labelled_index(n_negative=500) + second = _labelled_index(n_negative=500) + + first._apply_negative_ratio(50.0) + second._apply_negative_ratio(50.0) + + assert not torch.equal(first.negative_embeddings, second.negative_embeddings) + + def test_corpus_stays_aligned_with_embeddings(self): + """Every surviving embedding row still sits next to its own text.""" + index = _labelled_index(n_negative=40) + index._apply_negative_ratio(5.0) # 4 positives * 5 -> keep 20 + + assert index.negative_embeddings.shape[0] == 20 + assert len(index.negative_corpus) == 20 + + for row, text in zip(index.negative_embeddings, index.negative_corpus): + # Row i was filled with the value i, so it must pair with "neg i". + assert text == f"neg {int(row[0].item())}" + + def test_surviving_rows_keep_original_order(self): + """Sorting the kept indices preserves relative order, which aids debugging.""" + index = _labelled_index(n_negative=40) + index._apply_negative_ratio(5.0) + + values = [int(row[0].item()) for row in index.negative_embeddings] + assert values == sorted(values) + + def test_downsampling_without_corpus_still_works(self): + """An index with no corpus downsamples fine and stays corpus-free.""" + index = SentinelLocalIndex( + sentence_model=None, + positive_embeddings=torch.ones(4, 8), + negative_embeddings=torch.rand(40, 8), + ) + index._apply_negative_ratio(5.0, seed=7) + + assert index.negative_embeddings.shape[0] == 20 + assert index.negative_corpus is None + + def test_misaligned_corpus_is_dropped_with_warning(self, caplog): + """A corpus that does not match its embeddings is discarded, loudly. + + Falling back to row numbers is recoverable; naming the wrong sentence as the + reason for a score is not. + """ + import logging + + caplog.set_level(logging.WARNING) + + index = _labelled_index(n_negative=40) + index.negative_corpus = ["too", "few"] # deliberately misaligned + index._apply_negative_ratio(5.0, seed=7) + + assert index.negative_corpus is None + assert "does not match embedding rows" in caplog.text