mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
ca023be0c0
* Implement MuveraEmbedding * Add random generator parameter for reproducibility in MuveraEmbedding * Document random_seed parameter * Remove unnecessary module docstring from muvera_embedding.py * refactor: clean up constructor parameters and improve formatting in MuveraEmbedding * refactor: rename muvera_embedding.py to muvera.py and update related references * feat: enhance MuveraEmbedding with multi-vector model support and improve parameter defaults * feat: add embedding_size property to MuveraEmbedding * feat: update MuveraPostprocessor to use model description for embedding size and add Jupyter notebook for MUVERA usage * fix: fix types, doctest, rename variables, refactor (#545) * fix: fix types, doctest, rename variables, refactor * fix: fix python3.9 compatibility * fix: make get_output_dimension protected * Optimize muvera (#551) * vectorize operations * fix: fill empty clusters with dataset vectors * rollback get_output_dimension * fix: fix type hints * fix: review comments * tests: add tests --------- Co-authored-by: George <george.panchuk@qdrant.tech>
39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
import numpy as np
|
|
|
|
from fastembed import LateInteractionTextEmbedding
|
|
from fastembed.postprocess import Muvera
|
|
|
|
CANONICAL_VALUES = [-2.61810007e-04, 1.89005750e00, -2.32070747e00]
|
|
CANONICAL_QUERY_VALUES = [
|
|
-0.85783903,
|
|
1.1077204,
|
|
-0.09522747,
|
|
] # part of the values are zeros, should be compared with the result of nonzero mask
|
|
|
|
DIM = 128
|
|
K_SIM = 5
|
|
DIM_PROJ = 16
|
|
R_REPS = 20
|
|
|
|
|
|
def test_single_input():
|
|
model = LateInteractionTextEmbedding("colbert-ir/colbertv2.0", lazy_load=True)
|
|
random_generator = np.random.default_rng(42)
|
|
multivector = random_generator.random((10, 128))
|
|
|
|
for muvera in (
|
|
Muvera(dim=DIM, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS, random_seed=42),
|
|
Muvera.from_multivector_model(model, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS),
|
|
):
|
|
fde = muvera.process(multivector)
|
|
assert fde.shape[0] == muvera.embedding_size
|
|
assert np.allclose(fde[:3], CANONICAL_VALUES)
|
|
|
|
fde_doc = muvera.process_document(multivector)
|
|
assert fde_doc.shape[0] == muvera.embedding_size
|
|
assert np.allclose(fde, fde_doc)
|
|
|
|
fde_query = muvera.process_query(multivector)
|
|
assert fde_query.shape[0] == muvera.embedding_size
|
|
assert np.allclose(fde_query[np.nonzero(fde_query)][:3], CANONICAL_QUERY_VALUES)
|