fix some typing errors

Author: voorhs

src/autointent/_wrappers/embedder/base.py

@@ -38,7 +38,10 @@ def embed(
 
     @abstractmethod
     def embed(
-        self, utterances: list[str], task_type: TaskTypeEnum | None = None, return_tensors: bool = False
+        self,
+        utterances: list[str],
+        task_type: TaskTypeEnum | None = None,
+        return_tensors: bool = False,
     ) -> npt.NDArray[np.float32] | torch.Tensor:
         """Calculate embeddings for a list of utterances.
 

src/autointent/_wrappers/embedder/embedder.py

@@ -13,7 +13,6 @@
 import numpy as np
 import numpy.typing as npt
 import torch
-from typing_extensions import assert_never
 
 from autointent.configs import EmbedderFineTuningConfig, TaskTypeEnum
 from autointent.configs._embedder import EmbedderConfig, OpenaiEmbeddingConfig, SentenceTransformerEmbeddingConfig
@@ -54,10 +53,9 @@ def _load_model(self) -> BaseEmbeddingBackend:
         if isinstance(self.config, OpenaiEmbeddingConfig):
             return OpenaiEmbeddingBackend(self.config)
         # Check if it's exactly the abstract base config (not a subclass)
-        if type(self.config) is EmbedderConfig:
-            msg = f"Cannot instantiate abstract EmbedderConfig: {self.config.__repr__()}"
-            raise TypeError(msg)
-        assert_never(self.config)
+
+        msg = f"Cannot instantiate abstract EmbedderConfig: {self.config.__repr__()}"
+        raise TypeError(msg)
 
     def _get_hash(self) -> int:
         """Compute a hash value for the Embedder.
@@ -149,13 +147,9 @@ def load(cls, path: Path | str, override_config: EmbedderConfig | None = None) -
             instance._backend = SentenceTransformerEmbeddingBackend.load(backend_path)  # noqa: SLF001
         elif isinstance(config, OpenaiEmbeddingConfig):
             instance._backend = OpenaiEmbeddingBackend.load(backend_path)  # noqa: SLF001
-        # Check if it's exactly the abstract base config (not a subclass)
-        elif type(config) is EmbedderConfig:
-            # Handle abstract base config case
+        else:
             msg = f"Cannot load abstract EmbedderConfig: {config.__repr__()}"
             raise TypeError(msg)
-        else:
-            assert_never(config)
 
         return instance
 
@@ -170,7 +164,10 @@ def embed(
     ) -> npt.NDArray[np.float32]: ...
 
     def embed(
-        self, utterances: list[str], task_type: TaskTypeEnum | None = None, return_tensors: bool = False
+        self,
+        utterances: list[str],
+        task_type: TaskTypeEnum | None = None,
+        return_tensors: bool = False,
     ) -> npt.NDArray[np.float32] | torch.Tensor:
         """Calculate embeddings for a list of utterances.
 
@@ -182,7 +179,7 @@ def embed(
         Returns:
             A numpy array or PyTorch tensor of embeddings.
         """
-        return self._backend.embed(utterances, task_type, return_tensors=return_tensors)
+        return self._backend.embed(utterances=utterances, task_type=task_type, return_tensors=return_tensors)
 
     def similarity(
         self, embeddings1: npt.NDArray[np.float32], embeddings2: npt.NDArray[np.float32]

src/autointent/_wrappers/embedder/openai.py

@@ -2,13 +2,14 @@
 import logging
 from functools import partial
 from pathlib import Path
-from typing import Literal, overload
+from typing import Literal, TypedDict, cast, overload
 
 import aiometer
 import numpy as np
 import numpy.typing as npt
 import openai
 import torch
+from typing_extensions import NotRequired
 
 from autointent._hash import Hasher
 from autointent.configs import TaskTypeEnum
@@ -20,6 +21,12 @@
 logger = logging.getLogger(__name__)
 
 
+class EmbeddingsCreateKwargs(TypedDict):
+    input: list[str]
+    model: str
+    dimensions: NotRequired[int]
+
+
 class OpenaiEmbeddingBackend(BaseEmbeddingBackend):
     """OpenAI-based embedding backend implementation."""
 
@@ -30,9 +37,10 @@ def __init__(self, config: OpenaiEmbeddingConfig) -> None:
             config: Configuration for OpenAI embeddings.
         """
         self.config = config
-        self._client = None
-        self._async_client = None
-        self._event_loop = None
+        self._client: openai.OpenAI | None = None
+        self._async_client: openai.AsyncOpenAI | None = None
+        self._event_loop: asyncio.AbstractEventLoop | None = None
+
         if config.max_concurrent is not None:
             self._init_event_loop()
 
@@ -124,7 +132,7 @@ def embed(
             embeddings_path = get_embeddings_path(hasher.hexdigest())
             if embeddings_path.exists():
                 logger.debug("loading embeddings from %s", str(embeddings_path))
-                embeddings_np = np.load(embeddings_path).astype(np.float32)
+                embeddings_np = cast(npt.NDArray[np.float32], np.load(embeddings_path))
                 if return_tensors:
                     return torch.from_numpy(embeddings_np)
                 return embeddings_np
@@ -162,7 +170,7 @@ def _process_embeddings_sync(self, utterances: list[str]) -> np.ndarray:
             batch = utterances[i : i + self.config.batch_size]
 
             # Prepare API call parameters
-            kwargs = {
+            kwargs: EmbeddingsCreateKwargs = {
                 "input": batch,
                 "model": self.config.model_name,
             }
@@ -198,6 +206,9 @@ def _process_embeddings_async(self, utterances: list[str]) -> np.ndarray:
             max_at_once=self.config.max_concurrent,
             max_per_second=self.config.max_per_second,
         )
+        if self._event_loop is None:
+            msg = "Event loop is not initialized"
+            raise RuntimeError(msg)
         batch_results = self._event_loop.run_until_complete(task)
 
         # Flatten results
@@ -210,7 +221,7 @@ async def _process_batch_async(self, batch: list[str]) -> list[list[float]]:
         client = self._get_async_client()
 
         # Prepare API call parameters
-        kwargs = {
+        kwargs: EmbeddingsCreateKwargs = {
             "input": batch,
             "model": self.config.model_name,
         }
@@ -246,7 +257,7 @@ def similarity(
 
         # Calculate cosine similarity
         similarity_matrix = np.dot(normalized1, normalized2.T)
-        return similarity_matrix.astype(np.float32)
+        return cast(npt.NDArray[np.float32], similarity_matrix)
 
     def dump(self, path: Path) -> None:
         """Save the backend state to disk.

src/autointent/_wrappers/vector_index/vector_index.py

@@ -73,7 +73,7 @@ def add(self, texts: list[str], labels: ListOfLabels) -> None:
             msg = f"Texts and labels lengths mismatch: {len(texts)=} !] {len(labels)=}"
             raise ValueError(msg)
 
-        logger.debug("Adding embeddings to vector index %s", self.embedder.config.model_name)
+        logger.debug("Adding embeddings to vector index")
         embeddings = self.embedder.embed(texts, TaskTypeEnum.passage)
 
         if not hasattr(self, "index"):
@@ -85,7 +85,7 @@ def add(self, texts: list[str], labels: ListOfLabels) -> None:
 
     def clear_ram(self) -> None:
         """Clear the vector index from RAM."""
-        logger.debug("Clearing vector index %s from RAM", self.embedder.config.model_name)
+        logger.debug("Clearing vector index from RAM")
         self.embedder.clear_ram()
         self.index.clear_ram()
 

src/autointent/configs/__init__.py

@@ -6,6 +6,7 @@
     SentenceTransformerEmbeddingConfig,
     TaskTypeEnum,
     get_default_embedder_config,
+    initialize_embedder_config,
 )
 from ._inference_node import InferenceNodeConfig
 from ._optimization import DataConfig, HPOConfig, LoggingConfig
@@ -40,4 +41,5 @@
     "VocabConfig",
     "get_default_embedder_config",
     "get_default_vector_index_config",
+    "initialize_embedder_config",
 ]

src/autointent/configs/_embedder.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -98,5 +99,15 @@ class OpenaiEmbeddingConfig(EmbedderConfig):
     )
 
 
-def get_default_embedder_config() -> EmbedderConfig:
-    return SentenceTransformerEmbeddingConfig()
+def get_default_embedder_config(**kwargs: Any) -> EmbedderConfig:  # noqa: ANN401
+    return SentenceTransformerEmbeddingConfig(**kwargs)
+
+
+def initialize_embedder_config(values: dict[str, Any] | str | EmbedderConfig | None) -> EmbedderConfig:
+    if values is None:
+        return get_default_embedder_config()
+    if isinstance(values, EmbedderConfig):
+        return values.model_copy(deep=True)
+    if isinstance(values, str):
+        return get_default_embedder_config(model_name=values)
+    return get_default_embedder_config(**values)

src/autointent/modules/scoring/_catboost/catboost_scorer.py

@@ -11,7 +11,7 @@
 from pydantic import PositiveInt
 
 from autointent import Context, Embedder
-from autointent.configs import EmbedderConfig, TaskTypeEnum
+from autointent.configs import EmbedderConfig, TaskTypeEnum, initialize_embedder_config
 from autointent.custom_types import FloatFromZeroToOne, ListOfLabels
 from autointent.modules.base import BaseScorer
 
@@ -113,7 +113,7 @@ def __init__(
             msg = "Only catbooost text features will be used, `use_embedding_features` is ignored."
             logger.warning(msg)
 
-        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
+        self.embedder_config = initialize_embedder_config(embedder_config)
         self.loss_function = loss_function
         self.verbose = verbose
         self.catboost_kwargs = catboost_kwargs or {}

src/autointent/modules/scoring/_description/bi_encoder.py

@@ -7,7 +7,7 @@
 from pydantic import PositiveFloat
 
 from autointent import Context, Embedder
-from autointent.configs import EmbedderConfig, TaskTypeEnum
+from autointent.configs import EmbedderConfig, TaskTypeEnum, initialize_embedder_config
 
 from .base import BaseDescriptionScorer
 
@@ -64,7 +64,7 @@ def __init__(
         multilabel: bool = False,
     ) -> None:
         super().__init__(temperature=temperature, multilabel=multilabel)
-        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
+        self.embedder_config = initialize_embedder_config(embedder_config)
         self._embedder: Embedder | None = None
         self._description_vectors: NDArray[Any] | None = None
 

src/autointent/modules/scoring/_dnnc/dnnc.py

@@ -9,7 +9,13 @@
 from pydantic import PositiveInt
 
 from autointent import Context, Ranker, VectorIndex
-from autointent.configs import CrossEncoderConfig, EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
+from autointent.configs import (
+    CrossEncoderConfig,
+    EmbedderConfig,
+    VectorIndexConfig,
+    get_default_vector_index_config,
+    initialize_embedder_config,
+)
 from autointent.custom_types import Document, ListOfLabels
 from autointent.modules.base import BaseScorer
 
@@ -67,7 +73,7 @@ def __init__(
         vector_index_config: VectorIndexConfig | None = None,
     ) -> None:
         self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
-        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
+        self.embedder_config = initialize_embedder_config(embedder_config)
         self.k = k
         self.vector_index_config = vector_index_config or get_default_vector_index_config()
 

src/autointent/modules/scoring/_gcn/gcn_scorer.py

@@ -11,6 +11,7 @@
     EmbedderConfig,
     TaskTypeEnum,
     TorchTrainingConfig,
+    initialize_embedder_config,
 )
 from autointent.custom_types import ListOfLabels
 from autointent.modules.scoring._gcn.gcn_model import TextMLGCN
@@ -63,8 +64,8 @@ def __init__(  # noqa: PLR0913
     ) -> None:
         if gcn_hidden_dims is None:
             gcn_hidden_dims = [1024]
-        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
-        self.label_embedder_config = EmbedderConfig.from_search_config(label_embedder_config)
+        self.embedder_config = initialize_embedder_config(embedder_config)
+        self.label_embedder_config = initialize_embedder_config(label_embedder_config)
         self.gcn_hidden_dims = gcn_hidden_dims
         self.p_reweight = p_reweight
         self.tau_threshold = tau_threshold