new: update fastembed, add fastembed-gpu, add type alias for onnx pro…

…viders
qdrant · May 2, 2024 · bfa9bd0 · bfa9bd0
1 parent 2b86a09
commit bfa9bd0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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: