Skip to content

MPE: multi-prefix embeddings with chunked MaxSim retrieval - #229

Open
FarmersWrap wants to merge 61 commits into
texttron:mainfrom
FarmersWrap:mpe
Open

MPE: multi-prefix embeddings with chunked MaxSim retrieval#229
FarmersWrap wants to merge 61 commits into
texttron:mainfrom
FarmersWrap:mpe

Conversation

@FarmersWrap

Copy link
Copy Markdown

Summary

Adds multi-prefix embedding (MPE) support to the retriever: passages can be chunked (fixed size, random size range, independently encoded chunks, or pre-chunked input), and query–document scores are computed via MaxSim over chunk embeddings. Scripts under examples/mpe/ cover data preparation, training, and evaluation on MLDR-en, BrowseComp-Plus, and LongEmbed.

Relation to previous PRs

Supersedes #215 (and #211 before it). This is a fresh branch whose history never contained the decrypted BrowseComp-Plus topics-qrels/queries.tsv flagged in the #215 review — the file was purged via history rewrite, and a .gitignore entry now prevents it from being re-committed (the README documents regenerating it locally via decrypt_dataset.py --generate-tsv).

Status

The remaining review feedback from #215 is being addressed on this PR:

  • restore the forward-loss contract (route MaxSim through a chunk-aware DenseModel subclass so SPLADE/GradCache/subclass paths are untouched)
  • make the feature strictly opt-in — all chunk flags off means behavior identical to main
  • drop trust_remote_code, remove rich/debug leftovers
  • port the tests from 0121 maxsim support #211 as hermetic (no Hub downloads)
  • rebase onto current main (resolving the collator.py conflicts)

🤖 Generated with Claude Code

@FarmersWrap FarmersWrap mentioned this pull request Jul 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants