From ff2edd3c3c16434f74f2f14dc8fdbf78899371d6 Mon Sep 17 00:00:00 2001
From: michaelfeil <me@michaelfeil.eu>
Date: Mon, 17 Jun 2024 20:54:53 -0700
Subject: [PATCH] add pypi release

---
 .github/workflows/pypi_release.yaml           |  4 +--
 libs/infinity_emb/Dockerfile                  |  2 +-
 .../transformer/crossencoder/optimum.py       |  2 +-
 .../transformer/embedder/optimum.py           | 25 +++++++++----------
 .../infinity_emb/transformer/utils_optimum.py |  6 ++---
 5 files changed, 18 insertions(+), 21 deletions(-)

diff --git a/.github/workflows/pypi_release.yaml b/.github/workflows/pypi_release.yaml
index e0971e31..29cabb1c 100644
--- 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
\ No newline at end of file
+          echo PIP_VERSION=$(poetry run python -c "import infinity_emb;print(infinity_emb.__version__)") >> $GITHUB_OUTPUT
\ No newline at end of file
diff --git a/libs/infinity_emb/Dockerfile b/libs/infinity_emb/Dockerfile
index c35604cf..dfcf4985 100644
--- 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
index 87c75d78..937c84aa 100644
--- 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
index af248c8a..a3f06b0d 100644
--- 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
index 4021fd4b..f45553f9 100644
--- 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