fix: simplify CustomTextEmbedding __init__ parameter list

* fix(text): reduce parameter count in CustomTextEmbedding.__init__

Co-authored-by: n24q02m <135627235+n24q02m@users.noreply.github.com>

* fix(text): reduce parameter count and fix formatting

- Refactor CustomTextEmbedding.__init__ to use **kwargs.
- Implement CustomTextEmbeddingWorker for parallel processing support.
- Add unit test for CustomTextEmbeddingWorker.
- Fix ruff formatting in tests/test_custom_text_embedding.py.

Co-authored-by: n24q02m <135627235+n24q02m@users.noreply.github.com>

---------

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

Author: n24q02m

qwen3_embed/text/custom_text_embedding.py

@@ -1,19 +1,18 @@
-from collections.abc import Iterable, Sequence
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
 from numpy.typing import NDArray
 
-from qwen3_embed.common import OnnxProvider
 from qwen3_embed.common.model_description import (
     DenseModelDescription,
     PoolingType,
 )
 from qwen3_embed.common.onnx_model import OnnxOutputContext
-from qwen3_embed.common.types import Device, NumpyArray
+from qwen3_embed.common.types import NumpyArray
 from qwen3_embed.common.utils import last_token_pool, mean_pooling, normalize
-from qwen3_embed.text.onnx_embedding import OnnxTextEmbedding
+from qwen3_embed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 
 
 @dataclass(frozen=True)
@@ -29,26 +28,10 @@ class CustomTextEmbedding(OnnxTextEmbedding):
     def __init__(
         self,
         model_name: str,
-        cache_dir: str | None = None,
-        threads: int | None = None,
-        providers: Sequence[OnnxProvider] | None = None,
-        cuda: bool | Device = Device.AUTO,
-        device_ids: list[int] | None = None,
-        lazy_load: bool = False,
-        device_id: int | None = None,
-        specific_model_path: str | None = None,
         **kwargs: Any,
     ):
         super().__init__(
             model_name=model_name,
-            cache_dir=cache_dir,
-            threads=threads,
-            providers=providers,
-            cuda=cuda,
-            device_ids=device_ids,
-            lazy_load=lazy_load,
-            device_id=device_id,
-            specific_model_path=specific_model_path,
             **kwargs,
         )
         self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
@@ -58,6 +41,10 @@ def __init__(
     def _list_supported_models(cls) -> list[DenseModelDescription]:
         return cls.SUPPORTED_MODELS
 
+    @classmethod
+    def _get_worker_class(cls) -> type["CustomTextEmbeddingWorker"]:
+        return CustomTextEmbeddingWorker
+
     def _post_process_onnx_output(
         self, output: OnnxOutputContext, **kwargs: Any
     ) -> Iterable[NumpyArray]:
@@ -103,3 +90,18 @@ def add_model(
         cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
             pooling=pooling, normalization=normalization
         )
+
+
+class CustomTextEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> CustomTextEmbedding:
+        return CustomTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

tests/test_custom_text_embedding.py

@@ -201,6 +201,18 @@ def test_lazy_load_does_not_call_load(self, tmp_path: Path) -> None:
             CustomTextEmbedding(model_name=_MODEL_NAME, lazy_load=True)
         assert not called
 
+    def test_custom_worker_returns_correct_type(self, tmp_path: Path) -> None:
+        from qwen3_embed.text.custom_text_embedding import CustomTextEmbeddingWorker
+
+        _register(pooling=PoolingType.CLS)
+        with (
+            patch.object(CustomTextEmbedding, "download_model", return_value=tmp_path),
+            patch.object(CustomTextEmbedding, "load_onnx_model"),
+        ):
+            worker = CustomTextEmbeddingWorker(_MODEL_NAME, str(tmp_path))
+            emb = worker.init_embedding(_MODEL_NAME, str(tmp_path))
+        assert isinstance(emb, CustomTextEmbedding)
+
 
 # ===========================================================================
 # CustomTextEmbedding._pool (lines 66-89)