bump version to 0.6.1

* new: improve task setter in jina v3

* refactor

* new: add hf_token secret

* fix: cross platform env propagation

.github/workflows/python-tests.yml

@@ -42,5 +42,7 @@ jobs:
           poetry install --no-interaction --no-ansi --without dev,docs
 
       - name: Run pytest
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          poetry run pytest 
+          poetry run pytest

fastembed/parallel_processor.py

@@ -3,8 +3,6 @@ import os
 from collections import defaultdict
 from copy import deepcopy
 from enum import Enum
-from dataclasses import dataclass
-from multiprocessing import shared_memory, Manager, Lock
 from multiprocessing import Queue, get_context
 from multiprocessing.context import BaseContext
 from multiprocessing.process import BaseProcess
@@ -12,11 +10,6 @@ from multiprocessing.sharedctypes import Synchronized as BaseValue
 from queue import Empty
 from typing import Any, Iterable, Optional, Type
 
-import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common.types import NumpyArray
-
 
 # Single item should be processed in less than:
 processing_timeout = 10 * 60  # seconds
@@ -24,13 +17,6 @@ processing_timeout = 10 * 60  # seconds
 max_internal_batch_size = 200
 
 
-@dataclass
-class OnnxOutputContext:
-    model_output: NumpyArray
-    attention_mask: Optional[NDArray[np.int64]] = None
-    input_ids: Optional[NDArray[np.int64]] = None
-
-
 class QueueSignals(str, Enum):
     stop = "stop"
     confirm = "confirm"
@@ -46,54 +32,12 @@ class Worker:
         raise NotImplementedError()
 
 
-class SharedMemoryPool:
-    def __init__(self, lock: Lock):
-        self._lock = lock
-        self._pool: dict[str, tuple[shared_memory.SharedMemory, int, np.dtype]] = {}
-        self._free_buffers: list[str] = []
-
-    def allocate(self, size: int, dtype: np.dtype) -> tuple[shared_memory.SharedMemory, str]:
-        best_match = None
-        best_size = float("inf")
-        for buf_name in self._free_buffers:
-            shm, buf_size, buf_dtype = self._pool[buf_name]
-            # get best match for needed size
-            if buf_size >= size and buf_dtype == dtype and buf_size < best_size:
-                best_match = buf_name
-                best_size = buf_size
-            if best_match:
-                self._free_buffers.remove(best_match)
-                return self._pool[best_match][0], best_match
-            shm = shared_memory.SharedMemory(create=True, size=size)
-            self._pool[shm.name] = (shm, size, dtype)
-            return shm, shm.name
-        # if no match found, create new buffer
-        shm = shared_memory.SharedMemory(create=True, size=size)
-        self._pool[shm.name] = (shm, size, dtype)
-        return shm, shm.name
-
-    def release(self, name: str) -> None:
-        if name in self._pool and name not in self._free_buffers:
-            self._free_buffers.append(name)
-
-    def cleanup(self) -> None:
-        for shm, _, _ in self._pool.values():
-            shm.close()
-            shm.unlink()
-        self._pool.clear()
-        self._free_buffers.clear()
-
-    def __del__(self):
-        self.cleanup()
-
-
 def _worker(
     worker_class: Type[Worker],
     input_queue: Queue,
     output_queue: Queue,
     num_active_workers: BaseValue,
     worker_id: int,
-    shared_pool: SharedMemoryPool,
     kwargs: Optional[dict[str, Any]] = None,
 ) -> None:
     """
@@ -111,6 +55,7 @@ def _worker(
     try:
         worker = worker_class.start(**kwargs)
 
+        # Keep going until you get an item that's None.
         def input_queue_iterable() -> Iterable[Any]:
             while True:
                 item = input_queue.get()
@@ -119,24 +64,7 @@ def _worker(
                 yield item
 
         for processed_item in worker.process(input_queue_iterable()):
-            idx, output_context = processed_item
-            output_metadata = {}
-            for field in ["model_output", "attention_mask", "input_ids"]:
-                array = getattr(output_context, field, None)
-                if array is not None:
-                    shm, shm_name = shared_pool.allocate(array.nbytes, array.dtype)
-                    shm_array = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf)
-                    np.copyto(shm_array, array)
-                    output_metadata[field] = {
-                        "name": shm_name,
-                        "shape": array.shape,
-                        "dtype": array.dtype.str,
-                    }
-                    shm.close()
-            output_queue.put((idx, output_metadata))
-            for field in output_metadata:  # mark release to reuse
-                shared_pool.release(output_metadata[field]["name"])
-
+            output_queue.put(processed_item)
     except Exception as e:  # pylint: disable=broad-except
         logging.exception(e)
         output_queue.put(QueueSignals.error)
@@ -180,8 +108,6 @@ class ParallelWorkerPool:
         self.device_ids = device_ids
         self.cuda = cuda
         self.num_active_workers: Optional[BaseValue] = None
-        self.manager = Manager()
-        self.shared_pool = SharedMemoryPool(self.manager.Lock())
 
     def start(self, **kwargs: Any) -> None:
         self.input_queue = self.ctx.Queue(self.queue_size)
@@ -207,7 +133,6 @@ class ParallelWorkerPool:
                     self.output_queue,
                     self.num_active_workers,
                     worker_id,
-                    self.shared_pool,
                     worker_kwargs,
                 ),
             )
@@ -253,17 +178,7 @@ class ParallelWorkerPool:
                     if out_item == QueueSignals.error:
                         self.join_or_terminate()
                         raise RuntimeError("Thread unexpectedly terminated")
-
-                    idx, output_metadata = out_item
-                    output_arrays = {}
-                    for field, meta in output_metadata.items():
-                        shm = shared_memory.SharedMemory(name=meta["name"])
-                        array = np.ndarray(
-                            meta["shape"], dtype=meta["dtype"], buffer=shm.buf
-                        ).copy()
-                        output_arrays[field] = array
-                        shm.close()
-                    yield (idx, OnnxOutputContext(**output_arrays))
+                    yield out_item
                     read += 1
 
                 self.input_queue.put((idx, item))
@@ -278,18 +193,9 @@ class ParallelWorkerPool:
                 if out_item == QueueSignals.error:
                     self.join_or_terminate()
                     raise RuntimeError("Thread unexpectedly terminated")
-
-                idx, output_metadata = out_item
-                output_arrays = {}
-                for field, meta in output_metadata.items():
-                    shm = shared_memory.SharedMemory(name=meta["name"])
-                    array = np.ndarray(meta["shape"], dtype=meta["dtype"], buffer=shm.buf).copy()
-                    output_arrays[field] = array
-                    shm.close()
-                yield (idx, OnnxOutputContext(**output_arrays))
+                yield out_item
                 read += 1
         finally:
-            self.shared_pool.cleanup()
             assert self.input_queue is not None, "Input queue is None"
             assert self.output_queue is not None, "Output queue is None"
             self.join()
@@ -325,13 +231,11 @@ class ParallelWorkerPool:
             if process.is_alive():
                 process.terminate()
         self.processes.clear()
-        self.shared_pool.cleanup()
 
     def join(self) -> None:
         for process in self.processes:
             process.join()
         self.processes.clear()
-        self.shared_pool.cleanup()
 
     def __del__(self) -> None:
         """
@@ -346,4 +250,3 @@ class ParallelWorkerPool:
         for process in self.processes:
             if process.is_alive():
                 process.terminate()
-        self.shared_pool.cleanup()

fastembed/sparse/splade_pp.py

@@ -18,7 +18,7 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
         model_file="model.onnx",
     ),
     SparseModelDescription(
@@ -27,7 +27,7 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
         model_file="model.onnx",
     ),
 ]

fastembed/text/multitask_embedding.py

@@ -3,6 +3,7 @@ from typing import Any, Type, Iterable, Union, Optional
 
 import numpy as np
 
+from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import NumpyArray
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
@@ -44,9 +45,11 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, **kwargs: Any):
+    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
+        self.default_task_id: Union[Task, int] = (
+            task_id if task_id is not None else self.PASSAGE_TASK
+        )
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -57,9 +60,14 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+        self,
+        onnx_input: dict[str, NumpyArray],
+        task_id: Optional[Union[int, Task]] = None,
+        **kwargs: Any,
     ) -> dict[str, NumpyArray]:
-        onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
+        if task_id is None:
+            raise ValueError(f"task_id must be provided for JinaEmbeddingV3, got <{task_id}>")
+        onnx_input["task_id"] = np.array(task_id, dtype=np.int64)
         return onnx_input
 
     def embed(
@@ -67,20 +75,19 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        task_id: int = PASSAGE_TASK,
+        task_id: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
-        self.current_task_id = task_id
-        kwargs["task_id"] = task_id
-        yield from super().embed(documents, batch_size, parallel, **kwargs)
+        task_id = (
+            task_id if task_id is not None else self.default_task_id
+        )  # required for multiprocessing
+        yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
 
     def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.QUERY_TASK
-        yield from super().embed(query, **kwargs)
+        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.PASSAGE_TASK
-        yield from super().embed(texts, **kwargs)
+        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
 
 
 class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
@@ -90,11 +97,15 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         cache_dir: str,
         **kwargs: Any,
     ) -> JinaEmbeddingV3:
-        model = JinaEmbeddingV3(
+        return JinaEmbeddingV3(
             model_name=model_name,
             cache_dir=cache_dir,
             threads=1,
             **kwargs,
         )
-        model.current_task_id = kwargs["task_id"]
-        return model
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+        self.model: JinaEmbeddingV3  # mypy complaints `self.model` does not have `default_task_id`
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed(batch, task_id=self.model.default_task_id)
+            yield idx, onnx_output

fastembed/text/onnx_text_model.py

@@ -115,7 +115,7 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+                yield from self._post_process_onnx_output(self.onnx_embed(batch, **kwargs))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed"
-version = "0.6.0"
+version = "0.6.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"

tests/test_text_multitask_embeddings.py

@@ -109,9 +109,25 @@ def test_single_embedding():
 
             canonical_vector = task["vectors"]
             assert np.allclose(
-                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+                embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
             ), model_desc.model
 
+        classification_embeddings = list(model.embed(documents=docs, task_id=Task.CLASSIFICATION))
+        classification_embeddings = np.stack(classification_embeddings, axis=0)
+
+        assert classification_embeddings.shape == (len(docs), dim)
+
+        model = TextEmbedding(model_name=model_name, task_id=Task.CLASSIFICATION)
+        default_embeddings = list(model.embed(documents=docs))
+        default_embeddings = np.stack(default_embeddings, axis=0)
+
+        assert default_embeddings.shape == (len(docs), dim)
+
+        assert np.allclose(
+            classification_embeddings,
+            default_embeddings,
+            atol=1e-4,
+        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
@@ -140,7 +156,7 @@ def test_single_embedding_query():
 
         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
+            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -172,7 +188,7 @@ def test_single_embedding_passage():
 
         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
+            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -207,27 +223,6 @@ def test_parallel_processing(dim: int, model_name: str):
         delete_model_cache(model.model._model_dir)
 
 
-def test_task_assignment():
-    is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
-
-    if is_ci and not is_manual:
-        pytest.skip("Skipping in CI non-manual mode")
-
-    for model_desc in JinaEmbeddingV3._list_supported_models():
-        # todo: once we add more models, we should not test models >1GB size locally
-        model_name = model_desc.model
-
-        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: str):
     is_ci = os.getenv("CI")