mypy6

Author: SeBorgey

autointent/_dump_tools/unit_dumpers.py

@@ -12,7 +12,7 @@
 from peft import PeftModel
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
-from transformers import (  # type: ignore[attr-defined]
+from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
     PreTrainedModel,

autointent/context/data_handler/_stratification.py

@@ -13,7 +13,7 @@
 from numpy import typing as npt
 from sklearn.model_selection import train_test_split
 from skmultilearn.model_selection import IterativeStratification
-from transformers import set_seed  # type: ignore[attr-defined]
+from transformers import set_seed
 
 from autointent import Dataset
 from autointent.custom_types import LabelType

autointent/modules/scoring/_bert.py

@@ -9,7 +9,7 @@
 import torch
 from datasets import Dataset, DatasetDict
 from sklearn.model_selection import train_test_split
-from transformers import (  # type: ignore[attr-defined]
+from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
     DataCollatorWithPadding,

autointent/modules/scoring/_gcn/gcn_model.py

@@ -98,10 +98,10 @@ def set_correlation_matrix(self, train_labels: torch.Tensor) -> None:
         self.correlation_matrix.data.copy_(corr_matrix)
 
     def forward(self, bert_features: torch.Tensor, label_embeddings: torch.Tensor) -> torch.Tensor:
-        classifiers: torch.Tensor = label_embeddings
-        for i in range(len(self.gcn_layers)):
-            classifiers = self.gcn_layers[i](self.correlation_matrix, classifiers)
-            classifiers = self.activations[i](classifiers)
+        classifiers = label_embeddings
+        for gcn_layer, activation in zip(self.gcn_layers, self.activations, strict=True):
+            classifiers = gcn_layer(self.correlation_matrix, classifiers)
+            classifiers = activation(classifiers)
 
         return torch.matmul(bert_features, classifiers.T)