From cb513d9ce3446cfde7b3bdb02eacf2cd61e8c0a2 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:51:16 -0700 Subject: [PATCH 1/2] Persist the corpus and make index loading deterministic Three related defects in save/load, all of which undermine explainability or reproducibility. They ship together because fixing the first without the second would make things worse. 1. save() dropped the corpus. Only embeddings and config were written, so after any reload positive_corpus/negative_corpus were None and explanations reported the row number of a match instead of the matched sentence. Adds an optional corpus.json alongside the existing two files. 2. _apply_negative_ratio() reshuffled the embeddings without touching the corpus, so the two parallel lists came apart. Latent until now because no corpus was ever loaded - fix 1 would have activated it, turning "a useless number" into "a real but unrelated sentence presented as the reason for a score", which fools a reviewer in a way that obvious breakage does not. Both lists are now sliced together through one _take_rows helper. 3. torch.randperm ran without a generator, so every load() kept a different random subset of negatives. A saved index behaved like a slightly different model each time it was loaded: on the shipped example index the same text scored 0.028636 on one load and 0.011523 on the next. load() now takes an optional seed. Backward compatibility: corpus.json is optional, so indices saved by earlier versions load exactly as before. load_index() keeps its 3-tuple return value - the corpus is read by a separate load_corpus() rather than widening a function that external callers unpack positionally. Both new arguments are keyword-only or last-with-a-None-default, so positional callers are unaffected. Seeding uses a private torch.Generator rather than a global seed, so a caller's other randomness is left alone. Safety valves: a corpus whose length does not match its embeddings is refused at save time (ValueError) and discarded with a warning at load time. Degrading to row numbers is recoverable; naming the wrong sentence is not. Adds 13 tests covering the round trip, backward compatibility, seeded and unseeded downsampling, alignment after downsampling, and explanations naming real text after a reload. Co-authored-by: Cursor --- README.md | 18 ++++ src/sentinel/io/__init__.py | 4 +- src/sentinel/io/index_io.py | 98 +++++++++++++++++++- src/sentinel/sentinel_local_index.py | 116 +++++++++++++++++++++-- tests/test_io.py | 113 ++++++++++++++++++++++ tests/test_sentinel_local_index.py | 134 +++++++++++++++++++++++++++ 6 files changed, 472 insertions(+), 11 deletions(-) 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..fd11e28 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,12 @@ # Constants for file names CONFIG_FILE_NAME = "sentinel_local_index_config.json" EMBEDDINGS_FILE_NAME = "embeddings.safetensors" +CORPUS_FILE_NAME = "corpus.json" # Optional; absent in indices saved before corpus support # 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 +90,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 +132,18 @@ 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. + + 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 +186,20 @@ def save_index( ) as f_out: f_out.write(f_in.read()) + # 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. + if positive_corpus is not None or negative_corpus is not None: + 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 +263,48 @@ 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, and + both are None for indices saved without a corpus. + + 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..c04bb60 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,117 @@ 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(): + """Indices saved before corpus support have no corpus.json and must still load. + + This is the backward-compatibility guarantee: a missing file means "no corpus", + which is exactly the behaviour every existing index already has. + """ + 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), + ) + + # No corpus was passed, so the file should not have been created at all. + assert not os.path.exists(os.path.join(temp_dir, CORPUS_FILE_NAME)) + + # 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_deleting_corpus_file_degrades_gracefully(): + """Removing corpus.json from an index leaves the rest usable.""" + 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 From 4fe2c5b827ab68bf35b06a0614268fa05ff7c6d0 Mon Sep 17 00:00:00 2001 From: Wei Xiao <11197323+wxiao0421@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:49:15 -0700 Subject: [PATCH 2/2] Always write corpus.json so it cannot describe a previous save Skipping the corpus write when no corpus was supplied left an earlier save's corpus.json sitting next to freshly written embeddings. A stale corpus whose row count happens to match the new embeddings passes _corpus_if_aligned on load, so explanations named the wrong text with no warning at all. Index sizes are typically a fixed sample size, which makes a same-size rebuild the likely case rather than a corner one. Write the file unconditionally, storing nulls when there is no corpus. Callers keep both parameters optional, load_corpus already maps a JSON null to None exactly as it does a missing file, and an absent file still means the index was saved before corpus support. Saving without a corpus now clears any corpus at that path, which drops the implicit "keep the previous corpus" behaviour - only ever correct when the new embeddings are the same rows in the same order, and nothing enforced that. Co-authored-by: Cursor --- src/sentinel/io/index_io.py | 41 +++++++++++++++++++---------- tests/test_io.py | 51 ++++++++++++++++++++++++++++++++----- 2 files changed, 72 insertions(+), 20 deletions(-) diff --git a/src/sentinel/io/index_io.py b/src/sentinel/io/index_io.py index fd11e28..a587af0 100644 --- a/src/sentinel/io/index_io.py +++ b/src/sentinel/io/index_io.py @@ -37,7 +37,9 @@ # Constants for file names CONFIG_FILE_NAME = "sentinel_local_index_config.json" EMBEDDINGS_FILE_NAME = "embeddings.safetensors" -CORPUS_FILE_NAME = "corpus.json" # Optional; absent in indices saved before corpus support +# 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 @@ -136,6 +138,11 @@ def save_index( 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. """ @@ -186,19 +193,24 @@ 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. - if positive_corpus is not None or negative_corpus is not None: - 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, - ) + 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) @@ -279,8 +291,9 @@ def load_corpus( transport_params: Optional transport parameters for smart_open. Returns: - Tuple of (positive_corpus, negative_corpus). Either element may be None, and - both are None for indices saved without a corpus. + 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 diff --git a/tests/test_io.py b/tests/test_io.py index c04bb60..f6fd7e5 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -149,10 +149,10 @@ def test_corpus_round_trip(): def test_index_without_corpus_still_loads(): - """Indices saved before corpus support have no corpus.json and must still load. + """Saving with no corpus writes explicit nulls, and loading reports no corpus. - This is the backward-compatibility guarantee: a missing file means "no corpus", - which is exactly the behaviour every existing index already has. + 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( @@ -162,8 +162,10 @@ def test_index_without_corpus_still_loads(): negative_embeddings=torch.rand(4, 8), ) - # No corpus was passed, so the file should not have been created at all. - assert not os.path.exists(os.path.join(temp_dir, CORPUS_FILE_NAME)) + 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) @@ -171,8 +173,45 @@ def test_index_without_corpus_still_loads(): 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(): - """Removing corpus.json from an index leaves the rest usable.""" + """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,