Merge pull request #270 from michaelfeil/mf-pypi-fix

add pypi release
michaelfeil · Jun 18, 2024 · 498f433 · 498f433
2 parents cbb413b + ff2edd3
commit 498f433
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 21 deletions.
diff --git a/.github/workflows/pypi_release.yaml b/.github/workflows/pypi_release.yaml
@@ -92,9 +92,9 @@ jobs:
       - name: Install minimal dependencies and import
         shell: bash
         run: |
-          sleep 30
+          sleep 180
           poetry run pip install "infinity-emb[server]==${{ env.VERSION }}"
           poetry run infinity_emb --help
           poetry run python -c "import infinity_emb"
           # print version
-          echo version=$(poetry run python -c "import infinity_emb;print(infinity_emb.__version__)") >> $GITHUB_OUTPUT
+          echo PIP_VERSION=$(poetry run python -c "import infinity_emb;print(infinity_emb.__version__)") >> $GITHUB_OUTPUT
diff --git a/libs/infinity_emb/Dockerfile b/libs/infinity_emb/Dockerfile
@@ -1,5 +1,5 @@
 # Use the Python base image
-FROM nvidia/cuda:12.1.0-base-ubuntu22.04 AS base
+FROM nvidia/cuda:12.1.1-base-ubuntu22.04 AS base
 
 ENV PYTHONUNBUFFERED=1 \
     \

diff --git a/libs/infinity_emb/infinity_emb/transformer/crossencoder/optimum.py b/libs/infinity_emb/infinity_emb/transformer/crossencoder/optimum.py
@@ -18,7 +18,7 @@
             ORTModelForSequenceClassification,
         )
         from transformers import AutoConfig, AutoTokenizer  # type: ignore
-    except (ImportError, RuntimeError) as ex:
+    except (ImportError, RuntimeError, Exception) as ex:
         CHECK_ONNXRUNTIME.mark_dirty(ex)
 
 

diff --git a/libs/infinity_emb/infinity_emb/transformer/embedder/optimum.py b/libs/infinity_emb/infinity_emb/transformer/embedder/optimum.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 
+from infinity_emb._optional_imports import CHECK_ONNXRUNTIME, CHECK_TRANSFORMERS
 from infinity_emb.args import EngineArgs
 from infinity_emb.primitives import EmbeddingReturnType, PoolingMethod
 from infinity_emb.transformer.abstract import BaseEmbedder
@@ -16,24 +17,22 @@
     optimize_model,
 )
 
-try:
-    from optimum.onnxruntime import (  # type: ignore[import-untyped]
-        ORTModelForFeatureExtraction,
-    )
-    from transformers import AutoConfig, AutoTokenizer  # type: ignore[import-untyped]
+if CHECK_ONNXRUNTIME.is_available:
+    try:
+        from optimum.onnxruntime import (  # type: ignore[import-untyped]
+            ORTModelForFeatureExtraction,
+        )
 
-    OPTIMUM_AVAILABLE = True
-except (ImportError, RuntimeError):
-    OPTIMUM_AVAILABLE = False
+    except (ImportError, RuntimeError, Exception) as ex:
+        CHECK_ONNXRUNTIME.mark_dirty(ex)
+
+if CHECK_TRANSFORMERS.is_available:
+    from transformers import AutoConfig, AutoTokenizer  # type: ignore[import-untyped]
 
 
 class OptimumEmbedder(BaseEmbedder):
     def __init__(self, *, engine_args: EngineArgs):
-        if not OPTIMUM_AVAILABLE:
-            raise ImportError(
-                "optimum.onnxruntime is not installed."
-                "`pip install infinity_emb[optimum]`"
-            )
+        CHECK_ONNXRUNTIME.mark_required()
         provider = device_to_onnx(engine_args.device)
 
         onnx_file = get_onnx_files(

diff --git a/libs/infinity_emb/infinity_emb/transformer/utils_optimum.py b/libs/infinity_emb/infinity_emb/transformer/utils_optimum.py
@@ -13,10 +13,8 @@
     try:
         from optimum.onnxruntime import ORTOptimizer  # type: ignore
         from optimum.onnxruntime.configuration import OptimizationConfig  # type: ignore
-    except ImportError:
-        pass
-    except RuntimeError:
-        pass
+    except (ImportError, RuntimeError, Exception) as ex:
+        CHECK_ONNXRUNTIME.mark_dirty(ex)
 
 if CHECK_TORCH.is_available:
     import torch