Persist the corpus and make index loading deterministic

[wxiao0421]

Summary

Goal

To make a saved index a faithful, self-contained, reproducible artifact — so that loading one gives you back exactly what you saved, and gives you the same thing every time. Right now it fails both halves of that. Saving quietly loses information, and loading quietly adds randomness.

Fixes

Three related save/load bugs. They ship together because fixing the first without the second would make things worse.

1. save() dropped the corpus. Only the embeddings and config were written, so after any save() → load() round trip positive_corpus and negative_corpus were None, and explanations reported a row number instead of the matched sentence. Adds an optional corpus.json next to the existing two files.

2. _apply_negative_ratio() reshuffled the embeddings without the corpus, so the two parallel lists came apart. Latent on main because nothing loads a corpus today — fix 1 is what activates it, turning "a useless number" into "a real but unrelated sentence presented as the reason for a score". Both lists now go through a single _take_rows() helper.

3. torch.randperm ran without a generator, so every load() kept a different random subset of negatives and a saved index behaved like a slightly different model each time it was loaded. load() now takes an optional seed, using a private torch.Generator so the caller's other randomness is left alone.

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.

Safety valves

A corpus whose length doesn't match its embeddings is refused at save (ValueError) and discarded with a warning at load. Degrading to row numbers is recoverable; confidently naming the wrong sentence is not.

Test plan

13 new tests, 67 total (was 54): round trip, backward compatibility, seeded vs unseeded downsampling, alignment after downsampling, and explanations naming real text after a reload. flake8 output is unchanged from base and docs/check_docs_sync.py reports in sync. README documents corpus.json and seed=.

Base of #33, #34 and #35.

[wxiao0421]

Stack map (updated after #31 merged)

#31 has landed on main, and this stack has been rebased onto it. This PR is now the base of the remaining three.

main  (includes #31)
 └── #32  Persist the corpus + deterministic loading   (this PR)  [bug fixes]
      ├── #33  index.subsample()                                  [new method]
      │    └── #34  Grid-search index size / ratio axes           [new args]
      └── #35  SentinelLocalIndex.from_texts()                    [new classmethod]

Suggested merge order: #32 → #33 → #34, with #35 mergeable any time after #32.

#35 branches off #32 rather than the stack tip, because its only real dependency is the _take_rows helper. Verified that merging #35 into the tip of #34 produces no conflicts.

PR	Commits	Diff	What it does
#32	1	+472 / −11	corpus.json, corpus/embedding alignment, seeded loading
#33	1	+292	Resize an index without re-encoding
#34	1	+256 / −21	Sweep index size and ratio in run_grid_search
#35	2	+330 / −3	Build an index in one line

Test counts on each branch: 67, 82, 87 and 79 respectively, all passing, with flake8 output identical to each PR's base.

---

A note for whoever merges this. #31 was squash-merged, which collapsed its four commits into one new commit on main. Because the original SHAs are no longer ancestors of main, this PR temporarily re-proposed all of #31's work (5 commits, +2481/−707) and became unmergeable until the stack was rebased.

That will happen again to #33 and #35 each time a PR in this stack is squash-merged, and the remaining PRs will need the same rebase. If you'd prefer to avoid that churn, a merge commit or rebase merge on the rest of the stack keeps the commits as genuine ancestors of main and the downstream PRs update themselves.

[vcai4071]

Looks good overall. A few comments.

[vcai4071]

Do we want to require positive_corpus and/or negative_corpus? Otherwise, it seems we can overwrite the embeddings, but then if the corpuses aren't provided, then previously written corpuses would be used for explanations.

Alternatively, if they are not supplied, we can write null corpuses?

[wxiao0421]

Basically, positive_corpus and/or negative_corpus (text) are both option and it is added in this PR. We want to keep it this way and also work backward comparable.

Good catch, this is a real bug. Save embeddings + corpus to a path, then save new embeddings to that same path without a corpus, and corpus.json is left behind describing the old rows.

I will always write corpus.json and store nulls when the corpora aren't supplied, rather than requiring them.

[vcai4071]

Do we want to log an error and not write the corpuses if this is the case, rather than failing the entire run? Alternatively, we can at least write embeddings first, and then fail? Worried about calculating all the embeddings for nothing.

[wxiao0421]

The worry about wasted compute is fair, but I'd keep the raise. Two reasons.

This check can only fire on a caller bug. In every build path we document, the corpus is the same list that was encoded:

positive_embeddings = torch.tensor(model.encode(positive_examples, ...))
...
positive_corpus=positive_examples,

so the lengths match by construction, and a mismatch means the caller subsampled one side or passed the wrong variable. That's deterministic: it fails identically on rerun. Logging and skipping the corpus would instead write an expensive artifact that has permanently lost its explanations, and we'd likely only notice much later when production explanations show row numbers instead of text. That also undercuts the invariant the rest of the module works to keep (_corpus_if_aligned on load, _take_rows on subsetting).

On writing the embeddings first and then failing, I'd push back harder, because it will cause inconsistent embedding and corpus.json. It also rarely happens if run correctly.

[vcai4071]

We could potentially have the size check in the init, comparing the corpus to embeddings.

[wxiao0421]

Will check if we can add this in follow up PR.