fix mypy

Author: riapush

autointent/_callbacks/tensorboard.py

@@ -16,7 +16,7 @@ def __init__(self) -> None:
         Raises an ImportError if neither are installed.
         """
         try:
-            from torch.utils.tensorboard import SummaryWriter
+            from torch.utils.tensorboard import SummaryWriter # type: ignore[attr-defined]
 
             self.writer = SummaryWriter
         except ImportError:

autointent/_dump_tools.py

@@ -260,15 +260,15 @@ def load(  # noqa: C901, PLR0912, PLR0915
             elif child.name == Dumper.ptuning_models:
                 for model_dir in child.iterdir():
                     try:
-                        model = AutoModelForSequenceClassification.from_pretrained(model_dir / "base_model")
+                        model = AutoModelForSequenceClassification.from_pretrained(model_dir / "base_model") # type: ignore[no-untyped-call]
                         hf_models[model_dir.name] = PeftModel.from_pretrained(model, model_dir / "peft")
                     except Exception as e:  # noqa: PERF203
                         msg = f"Error loading PeftModel {model_dir.name}: {e}"
                         logger.exception(msg)
             elif child.name == Dumper.hf_models:
                 for model_dir in child.iterdir():
                     try:
-                        hf_models[model_dir.name] = AutoModelForSequenceClassification.from_pretrained(model_dir)
+                        hf_models[model_dir.name] = AutoModelForSequenceClassification.from_pretrained(model_dir) # type: ignore[no-untyped-call]
                     except Exception as e:  # noqa: PERF203
                         msg = f"Error loading HF model {model_dir.name}: {e}"
                         logger.exception(msg)

autointent/modules/scoring/_bert.py

@@ -76,7 +76,7 @@ def _initialize_model(self) -> Any:  # noqa: ANN401
         label2id = {i: i for i in range(self._n_classes)}
         id2label = {i: i for i in range(self._n_classes)}
 
-        return AutoModelForSequenceClassification.from_pretrained(
+        return AutoModelForSequenceClassification.from_pretrained( # type: ignore[no-untyped-call]
             self.classification_model_config.model_name,
             trust_remote_code=self.classification_model_config.trust_remote_code,
             num_labels=self._n_classes,

autointent/modules/scoring/_cnn/textcnn.py

@@ -68,7 +68,7 @@ def load(self, model_path: str) -> None:
         state_dict = torch.load(model_path)
         self.load_state_dict(state_dict)
 
-    def get_config(self) -> dict[str, int | list[int] | torch.Tensor | None]:
+    def get_config(self) -> dict[str, int | list[int] | torch.Tensor | float | None]:
         return {
             "vocab_size": self.vocab_size,
             "n_classes": self.n_classes,