debug: try out with markers

fastembed/common/onnx_model.py

@@ -1,7 +1,19 @@
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    Sequence,
+)
 
 import numpy as np
 import onnxruntime as ort
@@ -39,11 +51,21 @@ class OnnxModel(Generic[T]):
         model_dir: Path,
         model_file: str,
         threads: Optional[int],
+        providers: Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]] = None,
     ) -> None:
         model_path = model_dir / model_file
 
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
-        onnx_providers = ["CPUExecutionProvider"]
+
+        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        available_providers = ort.get_available_providers()
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
 
         so = ort.SessionOptions()
         so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

pyproject.toml

@@ -12,8 +12,16 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
 
 [tool.poetry.dependencies]
 python = ">=3.8.0,<3.13"
-onnx = "^1.15.0"
-onnxruntime = "^1.17.0"
+onnx = [
+    {version = "^1.15.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime = [
+    {version = "^1.17.0", optional = true, markers = "extra != 'gpu'"}
+]
+onnxruntime-gpu = [
+    { version = "^1.17.0", optional = true, python = "<3.13", markers = "extra == 'gpu'" }
+]
+
 tqdm = "^4.66"
 requests = "^2.31"
 tokenizers = "^0.15.1"
@@ -37,6 +45,8 @@ pillow = "^10.2.0"
 cairosvg = "^2.7.1"
 mknotebooks = "^0.8.0"
 
+[tool.poetry.extras]
+gpu = ["onnxruntime-gpu"]
 
 [build-system]
 requires = ["poetry-core"]