From bfa9bd0aa3018dcbc84f678d93e978f1e738abd3 Mon Sep 17 00:00:00 2001
From: George Panchuk <george.panchuk@qdrant.tech>
Date: Thu, 2 May 2024 20:58:23 +0200
Subject: [PATCH] new: update fastembed, add fastembed-gpu, add type alias for
 onnx providers

---
 pyproject.toml                          |  8 ++++----
 qdrant_client/async_qdrant_fastembed.py | 10 ++++++----
 qdrant_client/qdrant_fastembed.py       | 11 +++++++----
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 6bf65140..899e7940 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,10 +27,10 @@ grpcio-tools = ">=1.41.0"
 urllib3 = ">=1.26.14,<3"
 portalocker = "^2.7.0"
 fastembed = [
-    { version = "0.2.6", optional = true, python = "<3.13" }
+    { version = "0.2.7", optional = true, python = "<3.13" }
 ]
-onnxruntime-gpu = [
-    { version = "^1.17.0", optional = true, python = "<3.13" }
+fastembed-gpu = [
+    { version = "0.2.7", optional = true, python = "<3.13" }
 ]
 
 [tool.poetry.group.dev.dependencies]
@@ -61,7 +61,7 @@ types-protobuf = "^4.21.0.5"
 
 [tool.poetry.extras]
 fastembed = ["fastembed"]
-fastembed-gpu = ["fastembed", "onnxruntime-gpu"]
+fastembed-gpu = ["fastembed-gpu"]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
diff --git a/qdrant_client/async_qdrant_fastembed.py b/qdrant_client/async_qdrant_fastembed.py
index d44d958f..89b48d8b 100644
--- a/qdrant_client/async_qdrant_fastembed.py
+++ b/qdrant_client/async_qdrant_fastembed.py
@@ -22,9 +22,11 @@
 
 try:
     from fastembed import SparseTextEmbedding, TextEmbedding
+    from fastembed.common import OnnxProvider
 except ImportError:
     TextEmbedding = None
     SparseTextEmbedding = None
+    OnnxProvider = None
 SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = (
     {
         model["model"]: (model["dim"], models.Distance.COSINE)
@@ -73,7 +75,7 @@ def set_model(
         max_length: Optional[int] = None,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -115,7 +117,7 @@ def set_sparse_model(
         embedding_model_name: Optional[str],
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
     ) -> None:
         """
         Set sparse embedding model to use for hybrid search over documents in combination with dense embeddings.
@@ -168,7 +170,7 @@ def _get_or_init_model(
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> "TextEmbedding":
         if model_name in cls.embedding_models:
@@ -193,7 +195,7 @@ def _get_or_init_sparse_model(
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> "SparseTextEmbedding":
         if model_name in cls.sparse_embedding_models:
diff --git a/qdrant_client/qdrant_fastembed.py b/qdrant_client/qdrant_fastembed.py
index 810397a4..d45d3378 100644
--- a/qdrant_client/qdrant_fastembed.py
+++ b/qdrant_client/qdrant_fastembed.py
@@ -11,9 +11,12 @@
 
 try:
     from fastembed import SparseTextEmbedding, TextEmbedding
+    from fastembed.common import OnnxProvider
 except ImportError:
     TextEmbedding = None
     SparseTextEmbedding = None
+    OnnxProvider = None
+
 
 SUPPORTED_EMBEDDING_MODELS: Dict[str, Tuple[int, models.Distance]] = (
     {
@@ -67,7 +70,7 @@ def set_model(
         max_length: Optional[int] = None,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -112,7 +115,7 @@ def set_sparse_model(
         embedding_model_name: Optional[str],
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
     ) -> None:
         """
         Set sparse embedding model to use for hybrid search over documents in combination with dense embeddings.
@@ -170,7 +173,7 @@ def _get_or_init_model(
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> "TextEmbedding":
         if model_name in cls.embedding_models:
@@ -198,7 +201,7 @@ def _get_or_init_sparse_model(
         model_name: str,
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
-        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
+        providers: Optional[Sequence["OnnxProvider"]] = None,
         **kwargs: Any,
     ) -> "SparseTextEmbedding":
         if model_name in cls.sparse_embedding_models: