mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-28 21:51:07 -05:00
35 lines
997 B
Python
35 lines
997 B
Python
import os
|
|
|
|
import numpy as np
|
|
|
|
from fastembed.late_interaction.late_interaction_text_embedding import (
|
|
LateInteractionTextEmbedding,
|
|
)
|
|
from fastembed.common.pooling import LateInteractionPooler
|
|
|
|
from tests.utils import delete_model_cache
|
|
|
|
CANONICAL_COLUMN_VALUES = {
|
|
"colbert-ir/colbertv2.0": np.array(
|
|
[4.0727495e-03, -2.4026826e-03, -6.8204990e-04, -7.1383954e-05, 4.4963313e-03]
|
|
),
|
|
}
|
|
|
|
docs = ["Hello World"]
|
|
|
|
|
|
def test_batch_embedding():
|
|
is_ci = os.getenv("CI")
|
|
docs_to_embed = docs * 10
|
|
|
|
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
|
|
print("evaluating", model_name)
|
|
model = LateInteractionTextEmbedding(model_name=model_name)
|
|
pooler = LateInteractionPooler()
|
|
result = list(model.embed(docs_to_embed, batch_size=6))
|
|
pooled_result = pooler.pool(result)
|
|
assert np.allclose(pooled_result[0], expected_result, atol=2e-3)
|
|
|
|
if is_ci:
|
|
delete_model_cache(model.model._model_dir)
|