| """Sparse encoders for hybrid retrieval. |
| |
| Two channels — both produce sparse vectors in the format Qdrant expects |
| (parallel arrays of token-index ints and float weights). |
| |
| - SPLADE: learned sparse representation. Captures term importance + soft |
| term expansion. Strong on domain-specific corpora. |
| Model: prithivida/Splade_PP_en_v1 (SPLADE++, distilled). |
| (Replaces CLAUDE.md's naver/splade-cocondenser-ensembledistil — |
| same SPLADE family, fastembed-native, comparable quality.) |
| |
| - BM25: classic lexical bag-of-words with IDF + length normalization. |
| Stateless tokenization (Snowball stemmer) — no corpus fit needed |
| because fastembed's BM25 normalizes via fixed defaults; corpus- |
| relative IDF is applied at query time. Excellent baseline for |
| exact-term queries (regulatory codes, ticker symbols, fiscal |
| years, etc.). |
| |
| Output shape (per text): SparseVec = {"indices": [int, ...], "values": [float, ...]}. |
| Both are passed straight to qdrant_client.models.SparseVector at upsert time. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Iterable |
|
|
| from fastembed import SparseEmbedding, SparseTextEmbedding |
|
|
|
|
| SPLADE_MODEL = "prithivida/Splade_PP_en_v1" |
| BM25_MODEL = "Qdrant/bm25" |
|
|
|
|
| @dataclass |
| class SparseVec: |
| indices: list[int] |
| values: list[float] |
|
|
| def __len__(self) -> int: |
| return len(self.indices) |
|
|
|
|
| _MODEL_CACHE: dict[str, SparseTextEmbedding] = {} |
|
|
|
|
| def _get_model(model_name: str) -> SparseTextEmbedding: |
| if model_name not in _MODEL_CACHE: |
| _MODEL_CACHE[model_name] = SparseTextEmbedding(model_name=model_name) |
| return _MODEL_CACHE[model_name] |
|
|
|
|
| def _to_sparse_vec(emb: SparseEmbedding) -> SparseVec: |
| """fastembed returns numpy arrays; Qdrant wants Python lists.""" |
| return SparseVec( |
| indices=[int(i) for i in emb.indices.tolist()], |
| values=[float(v) for v in emb.values.tolist()], |
| ) |
|
|
|
|
| class SpladeEncoder: |
| """SPLADE++ encoder — same call shape for documents and queries.""" |
|
|
| def __init__(self, model_name: str = SPLADE_MODEL): |
| self.model_name = model_name |
| self.model = _get_model(model_name) |
|
|
| def encode(self, texts: list[str], *, batch_size: int = 16) -> list[SparseVec]: |
| return [_to_sparse_vec(e) for e in self.model.embed(texts, batch_size=batch_size)] |
|
|
| def encode_query(self, query: str) -> SparseVec: |
| return self.encode([query])[0] |
|
|
|
|
| class BM25Encoder: |
| """fastembed's BM25 — corpus-aware via Qdrant's native handling. |
| |
| Documents and queries are encoded with different methods because BM25 has |
| asymmetric weighting: doc weights include TF + length norm, query weights |
| are pure IDF. |
| """ |
|
|
| def __init__(self, model_name: str = BM25_MODEL): |
| self.model_name = model_name |
| self.model = _get_model(model_name) |
|
|
| def encode_documents(self, texts: list[str], *, batch_size: int = 32) -> list[SparseVec]: |
| return [_to_sparse_vec(e) for e in self.model.embed(texts, batch_size=batch_size)] |
|
|
| def encode_query(self, query: str) -> SparseVec: |
| |
| return _to_sparse_vec(next(self.model.query_embed(query))) |
|
|
