change CNNScorer similar to RNN one

Author: riapush

autointent/configs/_transformers.py

@@ -125,9 +125,9 @@ class CrossEncoderConfig(HFModelConfig):
 
 class CNNConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
+    device: str | None = Field(None, description="Torch notation for CPU or CUDA.")
     max_seq_length: int = Field(128, description="Maximum sequence length.")
     padding_idx: int = Field(0, description="Index used for padding.")
-    unknown_idx: int = Field(1, description="Index used for unknown.")
     batch_size: PositiveInt = Field(32, description="Batch size for model inference.")
 
     @classmethod

autointent/modules/scoring/_cnn/cnn.py

@@ -1,14 +1,10 @@
-"""CNNScorer class for scoring."""
-
-import re
-from collections import Counter
 from typing import Any
 
 import numpy as np
 import numpy.typing as npt
 import torch
 from torch import nn
-from torch.utils.data import DataLoader, TensorDataset
+from torch.optim import Adam
 
 from autointent import Context
 from autointent._callbacks import REPORTERS_NAMES
@@ -19,190 +15,200 @@
 
 
 class CNNScorer(BaseScorer):
-    """Convolutional Neural Network (CNN) scorer for intent classification."""
+    """Scorer based on CNN model for text classification."""
 
     name = "cnn"
-    supports_multilabel = True
     supports_multiclass = True
+    supports_multilabel = True
 
     def __init__(
         self,
-        num_train_epochs: int = 3,
-        learning_rate: float = 5e-5,
-        seed: int = 0,
-        report_to: REPORTERS_NAMES | None = None,  # type: ignore[valid-type]
         embed_dim: int = 128,
-        kernel_sizes: list[int] = [3, 4, 5], # noqa: B006
+        kernel_sizes: list[int] = [3, 4, 5],
         num_filters: int = 100,
         dropout: float = 0.1,
-        batch_size: int = 8,
         cnn_config: CNNConfig | str | dict[str, Any] | None = None,
+        num_train_epochs: int = 3,
+        batch_size: int = 8,
+        learning_rate: float = 5e-5,
+        seed: int = 0,
+        report_to: REPORTERS_NAMES | None = None,  # type: ignore  # noqa: PGH003
     ) -> None:
-        self.num_train_epochs = num_train_epochs
-        self.learning_rate = learning_rate
-        self.seed = seed
-        self.report_to = report_to
+        """Initialize the CNN scorer."""
         self.embed_dim = embed_dim
         self.kernel_sizes = kernel_sizes
         self.num_filters = num_filters
         self.dropout = dropout
         self.cnn_config = CNNConfig.from_search_config(cnn_config)
-
-        # Will be initialized during fit()
-        self._model: TextCNN | None = None
-        self._vocab: dict[str, int] | None = None
-        self._unk_token = "<UNK>"  # noqa: S105
-        self._pad_token = "<PAD>"  # noqa: S105
-        self._n_classes: int = 0
-        self._multilabel: bool = False
-        self._pad_idx = self.cnn_config.padding_idx
-        self._unk_idx = self.cnn_config.unknown_idx
-        self.batch_size = batch_size
-        self.max_seq_length = self.cnn_config.max_seq_length
+        self.num_train_epochs = num_train_epochs
+        self.batch_size = batch_size or self.cnn_config.batch_size
+        self.learning_rate = learning_rate
+        self.seed = seed
+        self.report_to = report_to
+        self._artifact = None
+        self._device = self.cnn_config.device or ("cuda" if torch.cuda.is_available() else "cpu")
 
     @classmethod
     def from_context(
         cls,
         context: Context,
+        embed_dim: int = 128,
+        kernel_sizes: list[int] = [3, 4, 5],
+        num_filters: int = 100,
+        dropout: float = 0.1,
+        cnn_config: CNNConfig | str | dict[str, Any] | None = None,
         num_train_epochs: int = 3,
         batch_size: int = 8,
         learning_rate: float = 5e-5,
         seed: int = 0,
-        embed_dim: int = 128,
-        kernel_sizes: list[int] = [3, 4, 5], # noqa: B006
-        num_filters: int = 100,
-        dropout: float = 0.1,
-        cnn_config: CNNConfig | str | dict[str, Any] | None = None
     ) -> "CNNScorer":
+        """Create a CNNScorer from context."""
+        report_to = context.logging_config.report_to
+
         return cls(
-            num_train_epochs=num_train_epochs,
-            batch_size=batch_size,
-            learning_rate=learning_rate,
-            seed=seed,
-            report_to=context.logging_config.report_to,
             embed_dim=embed_dim,
             kernel_sizes=kernel_sizes,
             num_filters=num_filters,
             dropout=dropout,
-            cnn_config=cnn_config
-        )
-
-    def get_implicit_initialization_params(self) -> dict[str, Any]:
-        return {"cnn_config": self.cnn_config.model_dump()}
-
-    def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
-        self._validate_task(labels)
-        self._multilabel = isinstance(labels[0], (list, np.ndarray)) # noqa: UP038
-
-        # Build vocabulary and tokenize
-        self._build_vocab(utterances)
-
-        # Convert text to padded indices
-        x = self._text_to_indices(utterances)
-        x_tensor = torch.tensor(x, dtype=torch.long)
-        y_tensor = torch.tensor(
-            labels, dtype=torch.long if not self._multilabel else torch.float
+            cnn_config=cnn_config,
+            num_train_epochs=num_train_epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            seed=seed,
+            report_to=report_to,
         )
 
-        # Initialize model
-        if self._vocab is None:
-            msg = "Vocabulary not built"
-            raise ValueError(msg)
-
+    def get_embedder_config(self) -> dict[str, Any]:
+        """Get the configuration of the embedder."""
+        config = self.cnn_config.model_dump()
+        config.update({
+            "embed_dim": self.embed_dim,
+            "kernel_sizes": self.kernel_sizes,
+            "num_filters": self.num_filters,
+            "dropout": self.dropout,
+        })
+        return config
+
+    def __initialize_model(self, vocab_size: int) -> None:
+        """Initialize the CNN model."""
         self._model = TextCNN(
-            vocab_size=len(self._vocab),
+            vocab_size=vocab_size,
             n_classes=self._n_classes,
             embed_dim=self.embed_dim,
             kernel_sizes=self.kernel_sizes,
             num_filters=self.num_filters,
             dropout=self.dropout,
-            padding_idx=self._pad_idx
+            padding_idx=self.cnn_config.padding_idx,
+            pretrained_embs=None,
         )
+        self._model.to(self.device)
+
+    def fit(
+        self,
+        utterances: list[str],
+        labels: ListOfLabels,
+    ) -> None:
+        """Fit the model to the given data."""
+        if hasattr(self, "_model"):
+            self.clear_cache()
+        self._validate_task(labels)
+        self._create_vocab(utterances)
+        self.__initialize_model(len(self._vocab))
+        x = self._texts_to_sequences(utterances)
+        y = torch.tensor(labels, dtype=torch.float) if self._multilabel else torch.tensor(labels, dtype=torch.long)
+        self._train_model(x, y)
+
+    def _create_vocab(self, utterances: list[str]) -> None:
+        """Create vocabulary from utterances."""
+        unique_words = set()
+        for text in utterances:
+            for word in text.lower().split():
+                unique_words.add(word)
+
+        self._vocab = {"<PAD>": 0, "<UNK>": 1}
+        for i, word in enumerate(unique_words):
+            self._vocab[word] = i + 2
+
+    def _texts_to_sequences(self, texts: list[str]) -> torch.Tensor:
+        """Convert texts to sequences using the vocabulary."""
+        sequences = [[self._vocab.get(word, self._vocab["<UNK>"]) for word in text.lower().split()] for text in texts]
+
+        max_len = min(max(len(seq) for seq in sequences), self.cnn_config.max_seq_length)
+        padded_sequences = [
+            seq[:max_len] if len(seq) > max_len else seq + [self._vocab["<PAD>"]] * (max_len - len(seq))
+            for seq in sequences
+        ]
+
+        return torch.tensor(padded_sequences, dtype=torch.long)
+
+    def _train_model(self, x: torch.Tensor, y: torch.Tensor) -> None:
+        """Train the model."""
+        self._model.train()
+        optimizer = Adam(self._model.parameters(), lr=self.learning_rate)
+
+        criterion = nn.BCEWithLogitsLoss() if self._multilabel else nn.CrossEntropyLoss()
+
+        x = x.to(self._device)
+        y = y.to(self._device)
+
+        dataset = torch.utils.data.TensorDataset(x, y)
+        dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
+
+        torch.manual_seed(self.seed)
+
+        for _epoch in range(self.num_train_epochs):
+            total_loss = 0
+            for batch_x, batch_y in dataloader:
+                optimizer.zero_grad()
+                outputs = self._model(batch_x)
+                loss = criterion(outputs, batch_y)
+                loss.backward()
+                optimizer.step()
+                total_loss += loss.item()
 
-        # Training
-        self._train_model(x_tensor, y_tensor)
+        self._model.eval()
 
     def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
-        if self._model is None:
-            msg = "Model not trained. Call fit() first."
-            raise ValueError(msg)
+        """Predict probabilities for utterances."""
+        if not hasattr(self, "_model") or not hasattr(self, "_vocab"):
+            msg = "Model is not trained. Call fit() first."
+            raise RuntimeError(msg)
 
-        x = self._text_to_indices(utterances)
-        x_tensor = torch.tensor(x, dtype=torch.long)
+        x = self._texts_to_sequences(utterances)
+        x = x.to(self.device)
 
         self._model.eval()
-        all_probs: list[npt.NDArray[Any]] = []
+        all_predictions = []
 
         with torch.no_grad():
-            for i in range(0, len(x_tensor), self.batch_size):
-                batch_x = x_tensor[i : i + self.batch_size]
+            for i in range(0, len(x), self.batch_size):
+                batch_x = x[i : i + self.batch_size]
                 outputs = self._model(batch_x)
+
                 if self._multilabel:
-                    probs = torch.sigmoid(outputs).cpu().numpy()
+                    batch_predictions = torch.sigmoid(outputs).cpu().numpy()
                 else:
-                    probs = torch.softmax(outputs, dim=1).cpu().numpy()
-                all_probs.append(probs)
-
-        return np.concatenate(all_probs, axis=0) if all_probs else np.array([])
-
-    def _build_vocab(self, utterances: list[str]) -> None:
-        """Build vocabulary from training utterances."""
-        word_counts: Counter[str] = Counter()
-        for utterance in utterances:
-            words = re.findall(r"\w+", utterance.lower())
-            word_counts.update(words)
-
-        # Create vocabulary with special tokens
-        self._vocab = {self._pad_token: self._pad_idx, self._unk_token: self._unk_idx}
-
-        # Convert Counter to list of (word, count) tuples sorted by frequency
-        sorted_words = word_counts.most_common()
-        for word, _ in sorted_words:
-            if word not in self._vocab:
-                self._vocab[word] = len(self._vocab)
-
-    def _text_to_indices(self, utterances: list[str]) -> list[list[int]]:
-        """Convert utterances to padded sequences of word indices."""
-        if self._vocab is None:
-            msg = "Vocabulary not built"
-            raise ValueError(msg)
-
-        sequences: list[list[int]] = []
-        for utterance in utterances:
-            words = re.findall(r"\w+", utterance.lower())
-            # Convert words to indices, using UNK for unknown words
-            seq = [self._vocab.get(word, self._unk_idx) for word in words]
-            # Truncate if too long
-            seq = seq[: self.max_seq_length]
-            # Pad if too short
-            seq = seq + [self._pad_idx] * (self.max_seq_length - len(seq))
-            sequences.append(seq)
-        return sequences
+                    batch_predictions = torch.softmax(outputs, dim=1).cpu().numpy()
 
-    def clear_cache(self) -> None:
-        self._model = None
-        torch.cuda.empty_cache()
+                all_predictions.append(batch_predictions)
 
-    def _train_model(self, x: torch.Tensor, y: torch.Tensor) -> None:
-        if self._model is None:
-            msg = "Model not initialized"
-            raise ValueError(msg)
+        return np.vstack(all_predictions) if all_predictions else np.array([])
 
-        dataset = TensorDataset(x, y)
-        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True)
+    def clear_cache(self) -> None:
+        """Clear model cache."""
+        if hasattr(self, "_model"):
+            del self._model
 
-        criterion = (
-            nn.CrossEntropyLoss() if not self._multilabel else nn.BCEWithLogitsLoss()
-        )
-        optimizer = torch.optim.Adam(self._model.parameters(), lr=self.learning_rate)
+    @property
+    def device(self) -> str:
+        """Get device used for model computations."""
+        return self._device
 
-        self._model.train()
-        for _ in range(self.num_train_epochs):
-            for batch_x, batch_y in dataloader:
-                optimizer.zero_grad()
-                outputs = self._model(batch_x)
-                loss = criterion(outputs, batch_y)
-                loss.backward()
-                optimizer.step()
+    @device.setter
+    def device(self, value: str) -> None:
+        """Set device for model computations."""
+        self._device = value
 
-        self._model.eval()
+    def get_implicit_initialization_params(self) -> dict[str, Any]:
+        """Return default params used in ``__init__`` method."""
+        return {"cnn_config": self.cnn_config.model_dump()}