new: add inference free splade#652
Conversation
📝 WalkthroughWalkthroughAdds an ONNX export experiment for the SPLADE document encoder with tokenizer metadata copying and Torch/ONNX parity validation. Introduces the Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
experiments/if_splade_to_onnx.py (1)
45-51: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winPin a single
revisionacross all downloads for reproducibility.
code_revisiononly pins the custom modeling code; the model weights (AutoModelForMaskedLM.from_pretrained), tokenizer, andidf.json(viahf_hub_download, lines 78-80) are all fetched from the default branch. If the upstream repo updates its vocab/idf files later, re-running this script could produce amodel.onnx+idf.json+ tokenizer combo that's inconsistent with what was validated byparity_check.♻️ Suggested fix: pin one revision for everything
+MODEL_REVISION = "main" # pin to a specific commit SHA once finalized + def export(model_id: str, output_dir: Path, opset: int = 14) -> Path: output_dir.mkdir(parents=True, exist_ok=True) model = AutoModelForMaskedLM.from_pretrained( - model_id, trust_remote_code=True, code_revision=CODE_REVISION + model_id, revision=MODEL_REVISION, trust_remote_code=True, code_revision=CODE_REVISION ) ... - tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer = AutoTokenizer.from_pretrained(model_id, revision=MODEL_REVISION) ... for file_name in TOKENIZER_FILES: - local_path = hf_hub_download(repo_id=model_id, filename=file_name) + local_path = hf_hub_download(repo_id=model_id, filename=file_name, revision=MODEL_REVISION)Also applies to: 78-83
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@experiments/if_splade_to_onnx.py` around lines 45 - 51, Pin a single revision for every Hugging Face download in the script: pass the shared revision to both AutoModelForMaskedLM.from_pretrained and AutoTokenizer.from_pretrained, and reuse it in the hf_hub_download call that retrieves idf.json. Keep code_revision separate for custom modeling code while ensuring the model, tokenizer, and IDF data all resolve to the same validated repository revision.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@fastembed/sparse/if_splade.py`:
- Around line 151-159: Update the eager initialization flow around
load_tokenizer and OnnxTextModel._load_onnx_model so the tokenizer loaded into
self.tokenizer and self.special_token_to_id is reused, or the second tokenizer
load is skipped. Preserve lazy_load behavior and ensure eager ONNX model
initialization still completes normally.
---
Nitpick comments:
In `@experiments/if_splade_to_onnx.py`:
- Around line 45-51: Pin a single revision for every Hugging Face download in
the script: pass the shared revision to both
AutoModelForMaskedLM.from_pretrained and AutoTokenizer.from_pretrained, and
reuse it in the hf_hub_download call that retrieves idf.json. Keep code_revision
separate for custom modeling code while ensuring the model, tokenizer, and IDF
data all resolve to the same validated repository revision.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2648919e-ba9d-4ec9-8748-e0d2f1ad6af8
📒 Files selected for processing (4)
experiments/if_splade_to_onnx.pyfastembed/sparse/if_splade.pyfastembed/sparse/sparse_text_embedding.pytests/test_sparse_embeddings.py
| # The tokenizer and the idf table are lightweight and are required for query embedding, | ||
| # which does not involve any model inference, so they are loaded eagerly, while | ||
| # `lazy_load` only defers the initialization of the onnx model | ||
| self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir) | ||
| self.special_tokens_ids: set[int] = set(self.special_token_to_id.values()) | ||
| self._token_id_to_idf = self._load_idf() | ||
|
|
||
| if not self.lazy_load: | ||
| self.load_onnx_model() |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether OnnxTextModel._load_onnx_model itself (re)loads the tokenizer.
fd -e py onnx_text_model.py --exec cat -n {}Repository: qdrant/fastembed
Length of output: 17592
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the Splade file and inspect the loader methods around the relevant lines.
ast-grep outline fastembed/sparse/if_splade.py --view expanded
echo "----"
cat -n fastembed/sparse/if_splade.py | sed -n '1,260p'Repository: qdrant/fastembed
Length of output: 13466
Avoid loading the tokenizer twice on eager init. load_tokenizer(model_dir=...) runs here, then OnnxTextModel._load_onnx_model() loads it again on the default lazy_load=False path, so startup pays the same parse/I/O cost twice. Reuse the first load or skip the second one.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastembed/sparse/if_splade.py` around lines 151 - 159, Update the eager
initialization flow around load_tokenizer and OnnxTextModel._load_onnx_model so
the tokenizer loaded into self.tokenizer and self.special_token_to_id is reused,
or the second tokenizer load is skipped. Preserve lazy_load behavior and ensure
eager ONNX model initialization still completes normally.
KShivendu
left a comment
KShivendu
left a comment
There was a problem hiding this comment.
Thanks for implementing this :D
2 participants
the only thing remaining is to upload the model to qdrant hf