Merge pull request #635 from michaelfeil/rw/allow-onnx-selection

expose additional params to cli to allow for better onnx selection

Author: wirthual

.github/workflows/test.yaml

@@ -47,7 +47,7 @@ jobs:
           android: true
           dotnet: true
           haskell: true
-          large-packages: false
+          large-packages: true
           docker-images: false
           swap-storage: false
       - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }}

libs/infinity_emb/infinity_emb/args.py

@@ -68,6 +68,8 @@ class EngineArgs:
     lengths_via_tokenize: bool = MANAGER.lengths_via_tokenize[0]
     embedding_dtype: EmbeddingDtype = EmbeddingDtype[MANAGER.embedding_dtype[0]]
     served_model_name: str = MANAGER.served_model_name[0]
+    onnx_disable_optimize: bool = MANAGER.onnx_disable_optimize[0]
+    onnx_do_not_prefer_quantized: bool = MANAGER.onnx_do_not_prefer_quantized[0]
 
     _loading_strategy: Optional[LoadingStrategy] = None
 
@@ -160,8 +162,10 @@ def from_env(cls) -> list["EngineArgs"]:
                 lengths_via_tokenize=lengths_via_tokenize,
                 embedding_dtype=embedding_dtype,
                 served_model_name=served_model_name,
+                onnx_disable_optimize=onnx_disable_optimize,
+                onnx_do_not_prefer_quantized=onnx_do_not_prefer_quantized
             )
-            for model_name_or_path, batch_size, revision, trust_remote_code, engine, model_warmup, device, compile, bettertransformer, dtype, pooling_method, lengths_via_tokenize, embedding_dtype, served_model_name in zip_longest(
+            for model_name_or_path, batch_size, revision, trust_remote_code, engine, model_warmup, device, compile, bettertransformer, dtype, pooling_method, lengths_via_tokenize, embedding_dtype, served_model_name,onnx_disable_optimize,onnx_do_not_prefer_quantized in zip_longest(
                 MANAGER.model_id,
                 MANAGER.batch_size,
                 MANAGER.revision,
@@ -176,5 +180,7 @@ def from_env(cls) -> list["EngineArgs"]:
                 MANAGER.lengths_via_tokenize,
                 MANAGER.embedding_dtype,
                 MANAGER.served_model_name,
+                MANAGER.onnx_disable_optimize,
+                MANAGER.onnx_do_not_prefer_quantized
             )
         ]

libs/infinity_emb/infinity_emb/cli.py

@@ -270,6 +270,14 @@ def v2(
             **_construct("proxy_root_path"),
             help="Proxy prefix for the application. See: https://fastapi.tiangolo.com/advanced/behind-a-proxy/",
         ),
+        onnx_disable_optimize: list[bool] = typer.Option(
+            **_construct("onnx_disable_optimize"),
+            help="Disable onnx optimization",
+        ),
+        onnx_do_not_prefer_quantized: list[bool] = typer.Option(
+            **_construct("onnx_do_not_prefer_quantized"),
+            help="Do not use quantized onnx models by default if available",
+        ),
     ):
         """Infinity API ♾️  cli v2. MIT License. Copyright (c) 2023-now Michael Feil \n
         \n
@@ -309,6 +317,8 @@ def v2(
         permissive_cors, bool: add permissive CORS headers to enable consumption from a browser. Defaults to False.
         api_key, str: optional Bearer token for authentication. Defaults to "", which disables authentication.
         proxy_root_path, str: optional Proxy prefix for the application. See: https://fastapi.tiangolo.com/advanced/behind-a-proxy/
+        onnx_disable_optimize, bool: disable onnx optimization
+        onnx_do_not_prefer_quantized, bool: do not prefer quantized onnx model if its available
         """
         logger.setLevel(log_level.to_int())
         device_id_typed = [DeviceID(d) for d in typer_option_resolve(device_id)]
@@ -330,6 +340,8 @@ def v2(
             compile=compile,
             bettertransformer=bettertransformer,
             served_model_name=served_model_name,
+            onnx_disable_optimize=onnx_disable_optimize,
+            onnx_do_not_prefer_quantized=onnx_do_not_prefer_quantized
         )
 
         engine_args = []

libs/infinity_emb/infinity_emb/env.py

@@ -260,5 +260,15 @@ def device_id(self):
     def embedding_dtype(self) -> list[str]:
         return self._typed_multiple("embedding_dtype", EmbeddingDtype)
 
-
+    @cached_property
+    def onnx_disable_optimize(self):
+        return self._to_bool_multiple(
+            self._optional_infinity_var_multiple("onnx_disable_optimize", default=["false"])
+        )
+    
+    @cached_property
+    def onnx_do_not_prefer_quantized(self):
+        return self._to_bool_multiple(
+            self._optional_infinity_var_multiple("onnx_do_not_prefer_quantized", default=["false"])
+        )
 MANAGER = __Infinity_EnvManager()

libs/infinity_emb/infinity_emb/transformer/classifier/optimum.py

@@ -2,7 +2,6 @@
 # Copyright (c) 2023-now michaelfeil
 
 import copy
-import os
 
 from infinity_emb._optional_imports import CHECK_ONNXRUNTIME, CHECK_TRANSFORMERS
 from infinity_emb.args import EngineArgs
@@ -36,7 +35,7 @@ def __init__(self, *, engine_args: EngineArgs):
             model_name_or_path=engine_args.model_name_or_path,
             revision=engine_args.revision,
             use_auth_token=True,
-            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()),
+            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()) and not engine_args.onnx_do_not_prefer_quantized,
         )
 
         model = optimize_model(
@@ -46,7 +45,7 @@ def __init__(self, *, engine_args: EngineArgs):
             trust_remote_code=engine_args.trust_remote_code,
             execution_provider=provider,
             file_name=onnx_file.as_posix(),
-            optimize_model=not os.environ.get("INFINITY_ONNX_DISABLE_OPTIMIZE", False),
+            optimize_model=not engine_args.onnx_disable_optimize
         )
         model.use_io_binding = False
 

libs/infinity_emb/infinity_emb/transformer/crossencoder/optimum.py

@@ -2,7 +2,6 @@
 # Copyright (c) 2023-now michaelfeil
 
 import copy
-import os
 
 import numpy as np
 
@@ -34,16 +33,14 @@ def __init__(self, *, engine_args: EngineArgs):
             model_name_or_path=engine_args.model_name_or_path,
             revision=engine_args.revision,
             use_auth_token=True,
-            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()),
+            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()) and not engine_args.onnx_do_not_prefer_quantized,
         )
 
         self.model = optimize_model(
             engine_args.model_name_or_path,
             execution_provider=provider,
             file_name=onnx_file.as_posix(),
-            optimize_model=not os.environ.get(
-                "INFINITY_ONNX_DISABLE_OPTIMIZE", False
-            ),  # TODO: make this env variable public
+            optimize_model=not engine_args.onnx_disable_optimize,
             model_class=ORTModelForSequenceClassification,
             revision=engine_args.revision,
             trust_remote_code=engine_args.trust_remote_code,

libs/infinity_emb/infinity_emb/transformer/embedder/optimum.py

@@ -2,7 +2,6 @@
 # Copyright (c) 2023-now michaelfeil
 
 import copy
-import os
 
 import numpy as np
 
@@ -42,7 +41,7 @@ def __init__(self, *, engine_args: EngineArgs):
             model_name_or_path=engine_args.model_name_or_path,
             revision=engine_args.revision,
             use_auth_token=True,
-            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()),
+            prefer_quantized=("cpu" in provider.lower() or "openvino" in provider.lower()) and not engine_args.onnx_do_not_prefer_quantized,
         )
 
         self.pooling = (
@@ -55,9 +54,7 @@ def __init__(self, *, engine_args: EngineArgs):
             trust_remote_code=engine_args.trust_remote_code,
             execution_provider=provider,
             file_name=onnx_file.as_posix(),
-            optimize_model=not os.environ.get(
-                "INFINITY_ONNX_DISABLE_OPTIMIZE", False
-            ),  # TODO: make this env variable public
+            optimize_model=not engine_args.onnx_disable_optimize,
             model_class=ORTModelForFeatureExtraction,
         )
         self.model.use_io_binding = False

libs/infinity_emb/infinity_emb/transformer/utils_optimum.py

@@ -144,7 +144,7 @@ def optimize_model(
     files_optimized = list(path_folder.glob(f"**/*{OPTIMIZED_SUFFIX}"))
 
     logger.info(f"files_optimized: {files_optimized}")
-    if files_optimized:
+    if files_optimized and optimize_model:
         file_optimized = files_optimized[-1]
         logger.info(f"Optimized model found at {file_optimized}, skipping optimization")
         return model_class.from_pretrained(