Persist the corpus and make index loading deterministic

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

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"]

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]:

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,22 +304,29 @@ 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.
 
         Args:
             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