mirror of
https://github.com/qdrant/fastembed.git
synced 2026-10-03 19:37:46 -05:00
PothanandGeorge Panchuk
44d0c4ef6d
* Fix MUVERA encoding for empty documents * Test MUVERA empty multivector encodings * fix: raise on empty vectors in muvera * Test ValueError for empty MUVERA inputs * replace getattr with explicit calls in tests Updated error messages for empty inputs in tests. * remove redundant fixture * remove unused import * pytest import is actually required * fix exception match message --------- Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
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)
|
|
|
|
|
|
def test_empty_multivectors_raise_value_error():
|
|
muvera = Muvera(dim=4, k_sim=2, dim_proj=2, r_reps=3)
|
|
empty = np.empty((0, 4))
|
|
|
|
with pytest.raises(ValueError, match="Cannot encode an empty multivector"):
|
|
muvera.process_document(empty)
|
|
|
|
with pytest.raises(ValueError, match="Cannot encode an empty multivector"):
|
|
muvera.process_query(empty)
|
|
|