mirror of
https://github.com/qdrant/fastembed.git
synced 2026-09-21 13:37:57 -05:00
62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
import numpy as np
|
|
|
|
from fastembed import (
|
|
TextEmbedding,
|
|
SparseTextEmbedding,
|
|
ImageEmbedding,
|
|
LateInteractionMultimodalEmbedding,
|
|
LateInteractionTextEmbedding,
|
|
)
|
|
from fastembed.common.utils import last_token_pooling
|
|
|
|
|
|
def test_text_list_supported_models():
|
|
for model_type in [
|
|
TextEmbedding,
|
|
SparseTextEmbedding,
|
|
ImageEmbedding,
|
|
LateInteractionMultimodalEmbedding,
|
|
LateInteractionTextEmbedding,
|
|
]:
|
|
supported_models = model_type.list_supported_models()
|
|
assert isinstance(supported_models, list)
|
|
description = supported_models[0]
|
|
assert isinstance(description, dict)
|
|
|
|
assert "model" in description and description["model"]
|
|
if model_type != SparseTextEmbedding:
|
|
assert "dim" in description and description["dim"]
|
|
assert "license" in description and description["license"]
|
|
assert "size_in_GB" in description and description["size_in_GB"]
|
|
assert "model_file" in description and description["model_file"]
|
|
assert "sources" in description and description["sources"]
|
|
assert "hf" in description["sources"] or "url" in description["sources"]
|
|
|
|
|
|
def test_last_token_pooling():
|
|
token_embeddings = np.array(
|
|
[
|
|
[[1.0, 1.0], [2.0, 2.0], [9.0, 9.0], [9.0, 9.0]], # 2 real tokens, then padding
|
|
[[3.0, 3.0], [4.0, 4.0], [5.0, 5.0], [6.0, 6.0]], # no padding
|
|
]
|
|
)
|
|
attention_mask = np.array([[1, 1, 0, 0], [1, 1, 1, 1]], dtype=np.int64)
|
|
|
|
pooled = last_token_pooling(token_embeddings, attention_mask)
|
|
|
|
assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
|
|
|
|
|
|
def test_last_token_pooling_with_left_padding():
|
|
token_embeddings = np.array(
|
|
[
|
|
[[9.0, 9.0], [9.0, 9.0], [1.0, 1.0], [2.0, 2.0]], # padding, then 2 real tokens
|
|
[[3.0, 3.0], [4.0, 4.0], [5.0, 5.0], [6.0, 6.0]], # no padding
|
|
]
|
|
)
|
|
attention_mask = np.array([[0, 0, 1, 1], [1, 1, 1, 1]], dtype=np.int64)
|
|
|
|
pooled = last_token_pooling(token_embeddings, attention_mask)
|
|
|
|
assert np.allclose(pooled, [[2.0, 2.0], [6.0, 6.0]])
|