fix typing errors

Author: voorhs

autointent/_dump_tools.py

@@ -10,7 +10,7 @@
 from peft import PeftModel
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
-from transformers import (
+from transformers import (  # type: ignore[attr-defined]
     AutoModelForSequenceClassification,
     AutoTokenizer,
     PreTrainedModel,
@@ -116,29 +116,31 @@ def dump(obj: Any, path: Path, exists_ok: bool = False, exclude: list[type[Any]]
                     if val._is_prompt_learning:  # noqa: SLF001
                         model_path = path / Dumper.peft_models / key
                         model_path.mkdir(parents=True, exist_ok=True)
-                        val.save_pretrained(model_path / "peft") # save peft config and prompt encoder
-                        val.base_model.save_pretrained(model_path / "base_model")  # save bert classifier
+                        # save peft config and prompt encoder
+                        val.save_pretrained(str(model_path / "peft"))
+                        # save bert classifier
+                        val.base_model.save_pretrained(model_path / "base_model")  # type: ignore[attr-defined]
                     else:
                         model_path = path / Dumper.hf_models / key
                         model_path.mkdir(parents=True, exist_ok=True)
                         merged_model: PreTrainedModel = val.merge_and_unload()
-                        merged_model.save_pretrained(model_path)
+                        merged_model.save_pretrained(model_path)  # type: ignore[attr-defined]
                 except Exception as e:
                     msg = f"Error dumping PeftModel {key}: {e}"
                     logger.exception(msg)
             elif isinstance(val, PreTrainedModel):
                 model_path = path / Dumper.hf_models / key
                 model_path.mkdir(parents=True, exist_ok=True)
                 try:
-                    val.save_pretrained(model_path)
+                    val.save_pretrained(model_path)  # type: ignore[attr-defined]
                 except Exception as e:
                     msg = f"Error dumping HF model {key}: {e}"
                     logger.exception(msg)
             elif isinstance(val, PreTrainedTokenizer | PreTrainedTokenizerFast):
                 tokenizer_path = path / Dumper.hf_tokenizers / key
                 tokenizer_path.mkdir(parents=True, exist_ok=True)
                 try:
-                    val.save_pretrained(tokenizer_path)
+                    val.save_pretrained(tokenizer_path)  # type: ignore[union-attr]
                 except Exception as e:
                     msg = f"Error dumping HF tokenizer {key}: {e}"
                     logger.exception(msg)

autointent/modules/scoring/_bert.py

@@ -11,9 +11,6 @@
     AutoModelForSequenceClassification,
     AutoTokenizer,
     DataCollatorWithPadding,
-    PreTrainedModel,
-    PreTrainedTokenizer,
-    PreTrainedTokenizerFast,
     Trainer,
     TrainingArguments,
 )
@@ -29,8 +26,8 @@ class BertScorer(BaseScorer):
     name = "bert"
     supports_multiclass = True
     supports_multilabel = True
-    _model: PreTrainedModel
-    _tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast
+    _model: Any  # transformers AutoModel factory returns Any
+    _tokenizer: Any  # transformers AutoTokenizer factory returns Any
 
     def __init__(
         self,
@@ -75,11 +72,11 @@ def from_context(
     def get_implicit_initialization_params(self) -> dict[str, Any]:
         return {"classification_model_config": self.classification_model_config.model_dump()}
 
-    def _initialize_model(self) -> None:
+    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)}
 
-        self._model = AutoModelForSequenceClassification.from_pretrained(
+        return AutoModelForSequenceClassification.from_pretrained(
             self.classification_model_config.model_name,
             trust_remote_code=self.classification_model_config.trust_remote_code,
             num_labels=self._n_classes,
@@ -99,7 +96,7 @@ def fit(
 
         self._tokenizer = AutoTokenizer.from_pretrained(self.classification_model_config.model_name)
 
-        self._initialize_model()
+        self._model = self._initialize_model()
 
         use_cpu = self.classification_model_config.device == "cpu"
 

autointent/modules/scoring/_lora/lora.py

@@ -77,7 +77,7 @@ def __init__(
             seed=seed,
             report_to=report_to,
         )
-        self._lora_config = LoraConfig(**lora_kwargs)  # type: ignore[arg-type]
+        self._lora_config = LoraConfig(**lora_kwargs)
 
     @classmethod
     def from_context(
@@ -102,6 +102,6 @@ def from_context(
             **lora_kwargs,
         )
 
-    def _initialize_model(self) -> None:
-        super()._initialize_model()
-        self._model = get_peft_model(self._model, self._lora_config)
+    def _initialize_model(self) -> Any:  # noqa: ANN401
+        model = super()._initialize_model()
+        return get_peft_model(model, self._lora_config)

autointent/modules/scoring/_ptuning/ptuning.py

@@ -69,7 +69,7 @@ def __init__(
             seed=seed,
             report_to=report_to,
         )
-        self._ptuning_config = PromptEncoderConfig(task_type="SEQ_CLS", **ptuning_kwargs)  # type: ignore[arg-type]
+        self._ptuning_config = PromptEncoderConfig(task_type="SEQ_CLS", **ptuning_kwargs)
 
     @classmethod
     def from_context(
@@ -108,7 +108,7 @@ def from_context(
             **ptuning_kwargs,
         )
 
-    def _initialize_model(self) -> None:
+    def _initialize_model(self) -> Any:  # noqa: ANN401
         """Initialize the model with P-tuning configuration."""
-        super()._initialize_model()
-        self._model = get_peft_model(self._model, self._ptuning_config)
+        model = super()._initialize_model()
+        return get_peft_model(model, self._ptuning_config)