diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 1e9978f7..0d36e118 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -1,11 +1,11 @@ from typing import Any, Iterable, Sequence, Type -from fastembed.common.types import NumpyArray, OnnxProvider, Device +from fastembed.common.model_description import DenseModelDescription, ModelSource from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.types import Device, NumpyArray, OnnxProvider from fastembed.common.utils import define_cache_dir, normalize from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker from fastembed.text.text_embedding_base import TextEmbeddingBase -from fastembed.common.model_description import DenseModelDescription, ModelSource supported_onnx_models: list[DenseModelDescription] = [ DenseModelDescription( @@ -180,6 +180,42 @@ sources=ModelSource(hf="jinaai/jina-clip-v1"), model_file="onnx/text_model.onnx", ), + DenseModelDescription( + model="minishlab/potion-base-8M", + dim=256, + description=( + "Text embeddings, Unimodal (text), English, 512 input tokens truncation, " + "Prefixes for queries/documents: not necessary, 2024 year." + ), + license="mit", + size_in_GB=0.030, + sources=ModelSource(hf="minishlab/potion-base-8m-onnx"), + model_file="model.onnx", + ), + DenseModelDescription( + model="minishlab/potion-retrieval-32M", + dim=512, + description=( + "Text embeddings, Unimodal (text), English, 512 input tokens truncation, " + "Prefixes for queries/documents: not necessary, 2025 year." + ), + license="mit", + size_in_GB=0.129, + sources=ModelSource(hf="minishlab/potion-retrieval-32m-onnx"), + model_file="model.onnx", + ), + DenseModelDescription( + model="minishlab/potion-multilingual-128M", + dim=256, + description=( + "Text embeddings, Unimodal (text), Multilingual, 512 input tokens truncation, " + "Prefixes for queries/documents: not necessary, 2025 year." + ), + license="mit", + size_in_GB=0.512, + sources=ModelSource(hf="minishlab/potion-multilingual-128m-onnx"), + model_file="model.onnx", + ), ] diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 8744b617..477655a8 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -82,6 +82,15 @@ "google/siglip2-base-patch16-224": np.array( [-0.01181389, 0.00737596, 0.01118064, 0.0103095, 0.3451049] ), + "minishlab/potion-base-8M": np.array( + [-0.03432461, -0.08020256, -0.14396408, 0.08480079, 0.01958815] + ), + "minishlab/potion-retrieval-32M": np.array( + [0.019733, -0.01530093, -0.08678473, 0.0229059, 0.04700558] + ), + "minishlab/potion-multilingual-128M": np.array( + [0.02366836, 0.02973341, 0.05140258, -0.00745248, -0.06740689] + ), } QWEN3_INSTRUCT_PREFIX = (