mypy1

Author: SeBorgey

autointent/_wrappers/base_torch_module.py

@@ -76,6 +76,9 @@ def __init__(
 
     def set_vocab(self, vocab: dict[str, Any]) -> None:
         """Save vocabulary into module's attributes and initialize embeddings matrix."""
+        if self.embed_dim is None:
+            msg = "embed_dim must be set to initialize embeddings"
+            raise ValueError(msg)
         self.vocab_config.vocab = vocab
         self.embedding = nn.Embedding(
             num_embeddings=len(self.vocab_config.vocab),

autointent/modules/scoring/_gcn/gcn_model.py

@@ -80,7 +80,7 @@ def create_correlation_matrix(
         adj_matrix = (conditional_prob > tau).float()
 
         adj_matrix_no_self_loop = adj_matrix - torch.eye(num_classes, device=adj_matrix.device)
-        sum_neighbors = adj_matrix_no_self_loop.sum(axis=1)
+        sum_neighbors = adj_matrix_no_self_loop.sum(dim=1)
 
         weights_p = p / sum_neighbors
         weights_p.nan_to_num_(0)

autointent/modules/scoring/_gcn/gcn_scorer.py

@@ -45,15 +45,16 @@ def __init__( # noqa: PLR0913
         self.gcn_hidden_dims = gcn_hidden_dims
         self.p_reweight = p_reweight
         self.tau_threshold = tau_threshold
-        self.torch_config = TorchTrainingConfig(
+        torch_config = TorchTrainingConfig(
             num_train_epochs=num_train_epochs,
             batch_size=batch_size,
             learning_rate=learning_rate,
             seed=seed,
         )
         if device is not None:
-            self.torch_config.device = device
-        self.early_stopping_config = EarlyStoppingConfig.from_search_config(early_stopping_config)
+            torch_config.device = device
+
+        super().__init__(torch_config=torch_config, early_stopping_config=early_stopping_config)
 
     @classmethod
     def from_context(  # noqa: PLR0913

autointent/modules/scoring/_torch/base_scorer.py

@@ -27,6 +27,14 @@ class BaseTorchTrainerScorer(BaseScorer):
     _best_model_weights = "best_model.pt"
     _model: BaseTorchModule
 
+    def __init__(
+        self,
+        torch_config: TorchTrainingConfig | dict[str, Any] | None = None,
+        early_stopping_config: EarlyStoppingConfig | dict[str, Any] | None = None,
+    ) -> None:
+        self.torch_config = TorchTrainingConfig.from_search_config(torch_config)
+        self.early_stopping_config = EarlyStoppingConfig.from_search_config(early_stopping_config)
+
     def _train_model(
         self,
         train_x: torch.Tensor,
@@ -170,9 +178,8 @@ def __init__(
         vocab_config: VocabConfig | dict[str, Any] | None = None,
         early_stopping_config: EarlyStoppingConfig | dict[str, Any] | None = None,
     ) -> None:
-        self.torch_config = TorchTrainingConfig.from_search_config(torch_config)
+        super().__init__(torch_config=torch_config, early_stopping_config=early_stopping_config)
         self.vocab_config = VocabConfig.from_search_config(vocab_config)
-        self.early_stopping_config = EarlyStoppingConfig.from_search_config(early_stopping_config)
 
     @abstractmethod
     def _init_model(self) -> BaseTorchModuleWithVocab: ...

autointent/modules/scoring/_torch/cnn_model.py

@@ -46,6 +46,7 @@ def __init__(
         self.kernel_sizes = kernel_sizes
         self.num_filters = num_filters
         self.dropout_rate = dropout
+        assert self.embed_dim is not None  # noqa: S101
 
         # Initialize other layers
         self.convs = nn.ModuleList(
@@ -70,6 +71,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.fc(dropped)  # type: ignore[no-any-return]
 
     def dump(self, path: Path) -> None:
+        assert self.embed_dim is not None  # noqa: S101
         metadata = TextCNNDumpMetadata(
             embed_dim=self.embed_dim,
             n_classes=self.n_classes,