typing

Author: SeBorgey

autointent/modules/__init__.py

@@ -49,7 +49,7 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
         MLKnnScorer,
         BertScorer,
         BERTLoRAScorer,
-        RNNScorer
+        RNNScorer,
     ]
 )
 

autointent/modules/scoring/__init__.py

@@ -19,4 +19,4 @@
     "RNNScorer",
     "RerankScorer",
     "SklearnScorer",
-    ]
+]

autointent/modules/scoring/_rnn.py

@@ -4,6 +4,7 @@
 import numpy.typing as npt
 import torch
 from torch import nn
+from torch.optim import Adam
 
 from autointent import Context
 from autointent._callbacks import REPORTERS_NAMES
@@ -63,15 +64,6 @@ def get_embedder_config(self) -> dict[str, Any]:
         """Get the configuration of the embedder."""
         return self.rnn_config.model_dump()
 
-    def _validate_task(self, labels: ListOfLabels) -> None:
-        """Validate the task type and set appropriate attributes."""
-        if isinstance(labels[0], list):
-            self._multilabel = True
-            self._n_classes = len(labels[0])
-        else:
-            self._multilabel = False
-            self._n_classes = max(labels) + 1
-
     def __initialize_model(self, vocab_size: int) -> None:
         """Initialize the RNN model."""
         self._model = SupervisedRNNClassifier(
@@ -128,7 +120,7 @@ def _texts_to_sequences(self, texts: list[str]) -> torch.Tensor:
     def _train_model(self, x: torch.Tensor, y: torch.Tensor) -> None:
         """Train the model."""
         self._model.train()
-        optimizer = torch.optim.Adam(self._model.parameters(), lr=self.learning_rate)
+        optimizer = Adam(self._model.parameters(), lr=self.learning_rate)
 
         criterion = nn.BCEWithLogitsLoss() if self._multilabel else nn.CrossEntropyLoss()
 
@@ -201,7 +193,7 @@ def __init__(
         super().__init__()
         if pretrained_embs is not None:
             _, embed_dim = pretrained_embs.shape
-            self.embedding = nn.Embedding.from_pretrained(pretrained_embs, freeze=True)
+            self.embedding = nn.Embedding.from_pretrained(pretrained_embs, freeze=True)  # type: ignore[no-untyped-call]
         else:
             self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
         self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers=n_layers, batch_first=True, dropout=dropout)

tests/modules/scoring/test_rnn.py

@@ -1,4 +1,3 @@
-
 import numpy as np
 import pytest
 
@@ -11,11 +10,7 @@ def test_rnn_prediction(dataset):
     """Test that the RNN model can fit and make predictions."""
     data_handler = DataHandler(dataset)
 
-    scorer = RNNScorer(
-        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1),
-        num_train_epochs=1,
-        batch_size=8
-    )
+    scorer = RNNScorer(rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1), num_train_epochs=1, batch_size=8)
 
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
@@ -52,11 +47,7 @@ def test_rnn_cache_clearing(dataset):
     """Test that the RNN model properly handles cache clearing."""
     data_handler = DataHandler(dataset)
 
-    scorer = RNNScorer(
-        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1),
-        num_train_epochs=1,
-        batch_size=8
-    )
+    scorer = RNNScorer(rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1), num_train_epochs=1, batch_size=8)
 
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
@@ -82,9 +73,7 @@ def test_rnn_device(dataset):
 
     # Force CPU
     scorer_cpu = RNNScorer(
-        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1, device="cpu"),
-        num_train_epochs=1,
-        batch_size=8
+        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1, device="cpu"), num_train_epochs=1, batch_size=8
     )
 
     scorer_cpu.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
@@ -97,9 +86,7 @@ def test_rnn_device(dataset):
 
     # Test with default device
     scorer_default = RNNScorer(
-        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1),
-        num_train_epochs=1,
-        batch_size=8
+        rnn_config=RNNConfig(embed_dim=64, hidden_dim=128, n_layers=1), num_train_epochs=1, batch_size=8
     )
 
     scorer_default.fit(data_handler.train_utterances(0), data_handler.train_labels(0))