fix: reduce code duplication in _load_onnx_model and add parallel execution support (#520)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: n24q02m

qwen3_embed/common/onnx_model.py

@@ -10,6 +10,7 @@
 from numpy.typing import NDArray
 from tokenizers import Tokenizer
 
+from qwen3_embed.common.preprocessor_utils import load_tokenizer
 from qwen3_embed.common.types import Device, NumpyArray, OnnxProvider
 from qwen3_embed.parallel_processor import Worker
 
@@ -48,6 +49,7 @@ def __init__(self) -> None:
         self.model: ort.InferenceSession | None = None
         self.model_input_names: set[str] | None = None
         self.tokenizer: Tokenizer | None = None
+        self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -108,9 +110,14 @@ def _validate_providers(
         return requested_provider_names
 
     def _create_session_options(
-        self, threads: int | None, extra_session_options: dict[str, Any] | None
+        self,
+        threads: int | None,
+        extra_session_options: dict[str, Any] | None,
+        parallel_execution: bool = False,
     ) -> Any:
         so = ort.SessionOptions()  # type: ignore[possibly-missing-attribute]
+        if parallel_execution:
+            so.execution_mode = ort.ExecutionMode.ORT_PARALLEL  # type: ignore[possibly-missing-attribute]
         so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  # type: ignore[possibly-missing-attribute]
         # Disable memory pattern optimization to prevent ORT from retaining
         # peak-sized buffers across inferences with varying sequence lengths.
@@ -124,32 +131,28 @@ def _create_session_options(
             self.add_extra_session_options(so, extra_session_options)
         return so
 
-    def _load_onnx_model(
+    def _instantiate_onnx_session(
         self,
-        model_dir: Path,
-        model_file: str,
+        model_path: Path,
         threads: int | None,
         providers: Sequence[OnnxProvider] | None = None,
         cuda: bool | Device = Device.AUTO,
         device_id: int | None = None,
+        parallel_execution: bool = False,
         extra_session_options: dict[str, Any] | None = None,
-    ) -> None:
-        model_path = model_dir / model_file
+    ) -> tuple[ort.InferenceSession, list[str]]:
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
         available_providers = ort.get_available_providers()  # type: ignore[possibly-missing-attribute]
 
         onnx_providers = self._determine_providers(providers, cuda, device_id, available_providers)
         requested_provider_names = self._validate_providers(onnx_providers, available_providers)
-        so = self._create_session_options(threads, extra_session_options)
+        so = self._create_session_options(threads, extra_session_options, parallel_execution)
 
-        self.model = ort.InferenceSession(
-            str(model_path), providers=onnx_providers, sess_options=so
-        )
-        self.model_input_names = {node.name for node in self.model.get_inputs()}
-        logger.info(f"ONNX session created with providers: {self.model.get_providers()}")
+        session = ort.InferenceSession(str(model_path), providers=onnx_providers, sess_options=so)
+        input_names = [node.name for node in session.get_inputs()]
+        logger.info(f"ONNX session created with providers: {session.get_providers()}")
         if "CUDAExecutionProvider" in requested_provider_names:
-            assert self.model is not None
-            current_providers = self.model.get_providers()
+            current_providers = session.get_providers()
             if "CUDAExecutionProvider" not in current_providers:
                 warnings.warn(
                     f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
@@ -158,6 +161,32 @@ def _load_onnx_model(
                     RuntimeWarning,
                     stacklevel=2,
                 )
+        return session, input_names
+
+    def _load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        parallel_execution: bool = False,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> tuple[ort.InferenceSession, list[str]]:
+        model_path = model_dir / model_file
+        self.model, input_names = self._instantiate_onnx_session(
+            model_path=model_path,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+            parallel_execution=parallel_execution,
+            extra_session_options=extra_session_options,
+        )
+        self.model_input_names = set(input_names)
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+        return self.model, input_names
 
     @classmethod
     def _select_exposed_session_options(cls, model_kwargs: dict[str, Any]) -> dict[str, Any]:

qwen3_embed/rerank/cross_encoder/onnx_text_model.py

@@ -1,7 +1,6 @@
 import os
 from collections.abc import Iterable, Sequence
 from multiprocessing import get_all_start_methods
-from pathlib import Path
 from typing import Any
 
 import numpy as np
@@ -13,7 +12,6 @@
     OnnxOutputContext,
     OnnxProvider,
 )
-from qwen3_embed.common.preprocessor_utils import load_tokenizer
 from qwen3_embed.common.types import Device, NumpyArray
 from qwen3_embed.common.utils import iter_batch
 from qwen3_embed.parallel_processor import ParallelWorkerPool
@@ -26,28 +24,6 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
     def _get_worker_class(cls) -> type["TextRerankerWorker"]:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _load_onnx_model(
-        self,
-        model_dir: Path,
-        model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
-    ) -> None:
-        super()._load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_file,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_id=device_id,
-            extra_session_options=extra_session_options,
-        )
-        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
-        assert self.tokenizer is not None
-
     def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
         assert self.tokenizer is not None
         return self.tokenizer.encode_batch(pairs)

qwen3_embed/text/onnx_text_model.py

@@ -1,15 +1,13 @@
 import os
 from collections.abc import Iterable, Sequence
 from multiprocessing import get_all_start_methods
-from pathlib import Path
 from typing import Any
 
 import numpy as np
 from numpy.typing import NDArray
-from tokenizers import Encoding, Tokenizer
+from tokenizers import Encoding
 
 from qwen3_embed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
-from qwen3_embed.common.preprocessor_utils import load_tokenizer
 from qwen3_embed.common.types import Device, NumpyArray, OnnxProvider
 from qwen3_embed.common.utils import iter_batch
 from qwen3_embed.parallel_processor import ParallelWorkerPool
@@ -36,8 +34,6 @@ def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) ->
 
     def __init__(self) -> None:
         super().__init__()
-        self.tokenizer: Tokenizer | None = None
-        self.special_token_to_id: dict[str, int] = {}
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], **kwargs: Any
@@ -47,27 +43,6 @@ def _preprocess_onnx_input(
         """
         return onnx_input
 
-    def _load_onnx_model(
-        self,
-        model_dir: Path,
-        model_file: str,
-        threads: int | None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_id: int | None = None,
-        extra_session_options: dict[str, Any] | None = None,
-    ) -> None:
-        super()._load_onnx_model(
-            model_dir=model_dir,
-            model_file=model_file,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_id=device_id,
-            extra_session_options=extra_session_options,
-        )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
-
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 

tests/test_cross_encoder_onnx.py

@@ -559,7 +559,7 @@ def test_wires_tokenizer_after_super(self, tmp_path: Path) -> None:
         with (
             patch("qwen3_embed.common.onnx_model.ort") as mock_ort,
             patch(
-                "qwen3_embed.rerank.cross_encoder.onnx_text_model.load_tokenizer",
+                "qwen3_embed.common.onnx_model.load_tokenizer",
                 return_value=(mock_tokenizer, {}),
             ),
         ):

tests/test_onnx_model.py

@@ -38,15 +38,20 @@ def model():
 
 @pytest.fixture
 def mock_ort():
-    with patch("qwen3_embed.common.onnx_model.ort") as mock:
+    with (
+        patch("qwen3_embed.common.onnx_model.ort") as mock,
+        patch("qwen3_embed.common.onnx_model.load_tokenizer", return_value=(MagicMock(), {})),
+    ):
         # Default behavior: CPU provider available, no CUDA
         mock.get_available_providers.return_value = ["CPUExecutionProvider"]
         mock.SessionOptions.return_value = MagicMock()
         mock.GraphOptimizationLevel.ORT_ENABLE_ALL = 99
+        mock.ExecutionMode.ORT_PARALLEL = 1
 
         # Default session mock
         session_mock = MagicMock()
         session_mock.get_providers.return_value = ["CPUExecutionProvider"]
+        session_mock.get_inputs.return_value = []
         mock.InferenceSession.return_value = session_mock
 
         yield mock
@@ -196,6 +201,16 @@ def test_load_threads(model: ConcreteOnnxModel, mock_ort):
     assert so.inter_op_num_threads == 4
 
 
+def test_load_parallel_execution(model: ConcreteOnnxModel, mock_ort):
+    """Test parallel execution configuration."""
+    mock_ort.get_available_providers.return_value = ["CPUExecutionProvider"]
+
+    model._load_onnx_model(Path("dummy"), "model.onnx", threads=None, parallel_execution=True)
+
+    so = mock_ort.SessionOptions.return_value
+    assert so.execution_mode == mock_ort.ExecutionMode.ORT_PARALLEL
+
+
 def test_load_extra_session_options(model: ConcreteOnnxModel, mock_ort):
     """Test extra session options."""
     mock_ort.get_available_providers.return_value = ["CPUExecutionProvider"]

tests/test_text_onnx_text_model.py

@@ -267,7 +267,7 @@ def test_wires_tokenizer_after_super(self, tmp_path: Path) -> None:
         with (
             patch("qwen3_embed.common.onnx_model.ort") as mock_ort,
             patch(
-                "qwen3_embed.text.onnx_text_model.load_tokenizer",
+                "qwen3_embed.common.onnx_model.load_tokenizer",
                 return_value=(mock_tokenizer, mock_special),
             ),
         ):