new: Added jina embedding v3 (#428)

* 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 <[email protected]>

* Update fastembed/text/multitask_embedding.py

Co-authored-by: George <[email protected]>

* 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 <[email protected]>

Author: hh-space-invader

NOTICE

@@ -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.
 

fastembed/text/multitask_embedding.py

@@ -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

fastembed/text/text_embedding.py

@@ -8,6 +8,7 @@
 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 @@ def embed(
             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)