mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
from typing import Any, Iterable, Type
|
|
|
|
import numpy as np
|
|
|
|
from fastembed.common.onnx_model import OnnxOutputContext
|
|
from fastembed.common.utils import normalize
|
|
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
|
|
from fastembed.text.onnx_text_model import TextEmbeddingWorker
|
|
from fastembed.text.pooled_embedding import PooledEmbedding
|
|
|
|
supported_pooled_normalized_models = [
|
|
{
|
|
"model": "sentence-transformers/all-MiniLM-L6-v2",
|
|
"dim": 384,
|
|
"description": "Text embeddings, Unimodal (text), English, 256 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
|
|
"license": "apache-2.0",
|
|
"size_in_GB": 0.09,
|
|
"sources": {
|
|
"url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
|
|
"hf": "qdrant/all-MiniLM-L6-v2-onnx",
|
|
},
|
|
"model_file": "model.onnx",
|
|
},
|
|
{
|
|
"model": "jinaai/jina-embeddings-v2-base-en",
|
|
"dim": 768,
|
|
"description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
|
|
"license": "apache-2.0",
|
|
"size_in_GB": 0.52,
|
|
"sources": {"hf": "xenova/jina-embeddings-v2-base-en"},
|
|
"model_file": "onnx/model.onnx",
|
|
},
|
|
{
|
|
"model": "jinaai/jina-embeddings-v2-small-en",
|
|
"dim": 512,
|
|
"description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
|
|
"license": "apache-2.0",
|
|
"size_in_GB": 0.12,
|
|
"sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
|
|
"model_file": "onnx/model.onnx",
|
|
},
|
|
{
|
|
"model": "jinaai/jina-embeddings-v2-base-de",
|
|
"dim": 768,
|
|
"description": "Text embeddings, Unimodal (text), Multilingual (German, English), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
|
|
"license": "apache-2.0",
|
|
"size_in_GB": 0.32,
|
|
"sources": {"hf": "jinaai/jina-embeddings-v2-base-de"},
|
|
"model_file": "onnx/model_fp16.onnx",
|
|
},
|
|
{
|
|
"model": "jinaai/jina-embeddings-v2-base-code",
|
|
"dim": 768,
|
|
"description": "Text embeddings, Unimodal (text), Multilingual (English, 30 programming languages), 8192 input tokens truncation, Prefixes for queries/documents: not necessary, 2024 year.",
|
|
"license": "apache-2.0",
|
|
"size_in_GB": 0.64,
|
|
"sources": {"hf": "jinaai/jina-embeddings-v2-base-code"},
|
|
"model_file": "onnx/model.onnx",
|
|
},
|
|
]
|
|
|
|
|
|
class PooledNormalizedEmbedding(PooledEmbedding):
|
|
@classmethod
|
|
def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
|
|
return PooledNormalizedEmbeddingWorker
|
|
|
|
@classmethod
|
|
def list_supported_models(cls) -> list[dict[str, Any]]:
|
|
"""Lists the supported models.
|
|
|
|
Returns:
|
|
list[dict[str, Any]]: A list of dictionaries containing the model information.
|
|
"""
|
|
return supported_pooled_normalized_models
|
|
|
|
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
|
|
if output.attention_mask is None:
|
|
raise ValueError("attention_mask must be provided for document post-processing")
|
|
|
|
embeddings = output.model_output
|
|
attn_mask = output.attention_mask
|
|
return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
|
|
|
|
|
|
class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
|
|
def init_embedding(
|
|
self,
|
|
model_name: str,
|
|
cache_dir: str,
|
|
**kwargs,
|
|
) -> OnnxTextEmbedding:
|
|
return PooledNormalizedEmbedding(
|
|
model_name=model_name,
|
|
cache_dir=cache_dir,
|
|
threads=1,
|
|
**kwargs,
|
|
)
|