upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

new: Added jina embedding v3 (#428)

Browse Source 
* new: Added jina embedding v3

* refactor: Changed dim to int value

* new: Updated notice

* new: Extended text embedding with query embed and passage embed

* fix: Fix lazy load in query and passage embed

* tests: Added test for multitask embeddings

* nit: Remove cache dir from tests

* tests: Updated tests

* improve: Improve task selection

* fix: Fix ci

* fix: Update fastembed/text/multitask_embedding.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* Update fastembed/text/multitask_embedding.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* fix: Pass task id using kwargs to parallel processor

* tests: Added test for task assignment

* prefer enums over ints

* tests: Added test for parallel

* improve: Updated model description

* fix: Fix ci

* fix: Fix ci

* refactor: Refactor query_embed and passage_embed

* tests: Added task propagation to parallel

* refactor: Set default task as retrieval passage

* chore: Update default task in tests

---------

Co-authored-by: George <george.panchuk@qdrant.tech>
This commit is contained in:
Hossam Hagag
2025-01-28 00:27:23 +02:00
committed by GitHub
parent 54f6cd9cbc
commit b05877de93
5 changed files with 388 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
NOTICE
View File

@@ -7,6 +7,8 @@ This distribution includes the following Jina AI models, each with its respectiv
- License: cc-by-nc-4.0
- jinaai/jina-reranker-v2-base-multilingual
- License: cc-by-nc-4.0
- jinaai/jina-embeddings-v3
- License: cc-by-nc-4.0
These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.

98
fastembed/text/multitask_embedding.py Normal file
View File

@@ -0,0 +1,98 @@
from enum import Enum
from typing import Any, Type, Iterable, Union, Optional
import numpy as np
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_multitask_models = [
{
"model": "jinaai/jina-embeddings-v3",
"dim": 1024,
"tasks": {
"retrieval.query": 0,
"retrieval.passage": 1,
"separation": 2,
"classification": 3,
"text-matching": 4,
},
"description": "Multi-task unimodal (text) embedding model, multi-lingual (~100), 1024 tokens truncation, and 8192 sequence length. Prefixes for queries/documents: not necessary, 2024 year.",
"license": "cc-by-nc-4.0",
"size_in_GB": 2.29,
"sources": {
"hf": "jinaai/jina-embeddings-v3",
},
"model_file": "onnx/model.onnx",
"additional_files": ["onnx/model.onnx_data"],
},
]
class Task(int, Enum):
RETRIEVAL_QUERY = 0
RETRIEVAL_PASSAGE = 1
SEPARATION = 2
CLASSIFICATION = 3
TEXT_MATCHING = 4
class JinaEmbeddingV3(PooledNormalizedEmbedding):
PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
QUERY_TASK = Task.RETRIEVAL_QUERY
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._current_task_id = self.PASSAGE_TASK
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
return JinaEmbeddingV3Worker
@classmethod
def list_supported_models(cls) -> list[dict[str, Any]]:
return supported_multitask_models
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
onnx_input["task_id"] = np.array(self._current_task_id, dtype=np.int64)
return onnx_input
def embed(
self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
task_id: int = PASSAGE_TASK,
**kwargs,
) -> Iterable[np.ndarray]:
self._current_task_id = task_id
kwargs["task_id"] = task_id
yield from super().embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.QUERY_TASK
yield from super().embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.PASSAGE_TASK
yield from super().embed(texts, **kwargs)
class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> JinaEmbeddingV3:
model = JinaEmbeddingV3(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)
model._current_task_id = kwargs["task_id"]
return model

29
fastembed/text/text_embedding.py
View File

@@ -8,6 +8,7 @@ from fastembed.text.clip_embedding import CLIPOnnxEmbedding
from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.pooled_embedding import PooledEmbedding
from fastembed.text.multitask_embedding import JinaEmbeddingV3
from fastembed.text.onnx_embedding import OnnxTextEmbedding
from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -19,6 +20,7 @@ class TextEmbedding(TextEmbeddingBase):
CLIPOnnxEmbedding,
PooledNormalizedEmbedding,
PooledEmbedding,
JinaEmbeddingV3,
]
@classmethod
@@ -113,3 +115,30 @@ class TextEmbedding(TextEmbeddingBase):
List of embeddings, one per document
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
"""
Embeds queries
Args:
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[np.ndarray]: The embeddings.
"""
# This is model-specific, so that different models can have specialized implementations
yield from self.model.query_embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
Args:
texts (Iterable[str]): The list of texts to embed.
**kwargs: Additional keyword argument to pass to the embed method.
Yields:
Iterable[SparseEmbedding]: The sparse embeddings.
"""
# This is model-specific, so that different models can have specialized implementations
yield from self.model.passage_embed(texts, **kwargs)

253
tests/test_text_multitask_embeddings.py Normal file
View File

@@ -0,0 +1,253 @@
import os
import numpy as np
import pytest
from fastembed import TextEmbedding
from fastembed.text.multitask_embedding import Task
from tests.utils import delete_model_cache
CANONICAL_VECTOR_VALUES = {
"jinaai/jina-embeddings-v3": [
{
"task_id": Task.RETRIEVAL_QUERY,
"vectors": np.array(
[
[0.0623, -0.0402, 0.1706, -0.0143, 0.0617],
[-0.1064, -0.0733, 0.0353, 0.0096, 0.0667],
]
),
},
{
"task_id": Task.RETRIEVAL_PASSAGE,
"vectors": np.array(
[
[0.0513, -0.0247, 0.1751, -0.0075, 0.0679],
[-0.0987, -0.0786, 0.09, 0.0087, 0.0577],
]
),
},
{
"task_id": Task.SEPARATION,
"vectors": np.array(
[
[0.094, -0.1065, 0.1305, 0.0547, 0.0556],
[0.0315, -0.1468, 0.065, 0.0568, 0.0546],
]
),
},
{
"task_id": Task.CLASSIFICATION,
"vectors": np.array(
[
[0.0606, -0.0877, 0.1384, 0.0065, 0.0722],
[-0.0502, -0.119, 0.032, 0.0514, 0.0689],
]
),
},
{
"task_id": Task.TEXT_MATCHING,
"vectors": np.array(
[
[0.0911, -0.0341, 0.1305, -0.026, 0.0576],
[-0.1432, -0.05, 0.0133, 0.0464, 0.0789],
]
),
},
]
}
docs = ["Hello World", "Follow the white rabbit."]
def test_batch_embedding():
is_ci = os.getenv("CI")
docs_to_embed = docs * 10
default_task = Task.RETRIEVAL_PASSAGE
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} default task")
embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs_to_embed), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding():
is_ci = os.getenv("CI")
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
for task in CANONICAL_VECTOR_VALUES[model_name]:
print(f"evaluating {model_name} task_id: {task['task_id']}")
embeddings = list(model.embed(documents=docs, task_id=task["task_id"]))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = task["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding_query():
is_ci = os.getenv("CI")
task_id = Task.RETRIEVAL_QUERY
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} query_embed task_id: {task_id}")
embeddings = list(model.query_embed(query=docs))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding_passage():
is_ci = os.getenv("CI")
task_id = Task.RETRIEVAL_PASSAGE
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} passage_embed task_id: {task_id}")
embeddings = list(model.passage_embed(texts=docs))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_parallel_processing():
is_ci = os.getenv("CI")
docs = ["Hello World", "Follow the white rabbit."] * 10
model_name = "jinaai/jina-embeddings-v3"
dim = 1024
model = TextEmbedding(model_name=model_name)
task_id = Task.SEPARATION
embeddings_1 = list(model.embed(docs, batch_size=10, parallel=None, task_id=task_id))
embeddings_1 = np.stack(embeddings_1, axis=0)
embeddings_2 = list(model.embed(docs, batch_size=10, parallel=1, task_id=task_id))
embeddings_2 = np.stack(embeddings_2, axis=0)
assert embeddings_1.shape[0] == len(docs) and embeddings_1.shape[-1] == dim
assert np.allclose(embeddings_1, embeddings_2, atol=1e-4)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(embeddings_2[:2, : canonical_vector.shape[1]], canonical_vector, atol=1e-4)
if is_ci:
delete_model_cache(model.model._model_dir)
def test_task_assignment():
is_ci = os.getenv("CI")
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
for i, task_id in enumerate(Task):
_ = list(model.embed(documents=docs, batch_size=1, task_id=i))
assert model.model._current_task_id == task_id
if is_ci:
delete_model_cache(model.model._model_dir)
@pytest.mark.parametrize(
"model_name",
["jinaai/jina-embeddings-v3"],
)
def test_lazy_load(model_name):
is_ci = os.getenv("CI")
model = TextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
list(model.embed(docs))
assert hasattr(model.model, "model")
if is_ci:
delete_model_cache(model.model._model_dir)

8
tests/test_text_onnx_embeddings.py
View File

@@ -69,14 +69,18 @@ CANONICAL_VECTOR_VALUES = {
"jinaai/jina-clip-v1": np.array([-0.0862, -0.0101, -0.0056, 0.0375, -0.0472]),
}
MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
def test_embedding():
is_ci = os.getenv("CI")
is_mac = platform.system() == "Darwin"
for model_desc in TextEmbedding.list_supported_models():
if (not is_ci and model_desc["size_in_GB"] > 1) or (
is_mac and model_desc["model"] == "nomic-ai/nomic-embed-text-v1.5-Q"
if (
(not is_ci and model_desc["size_in_GB"] > 1)
or model_desc["model"] in MULTI_TASK_MODELS
or (is_mac and model_desc["model"] == "nomic-ai/nomic-embed-text-v1.5-Q")
):
continue