added wrapper for skorch models to not need to specify classes on init, but on fit

pyproject.toml

 [tool.poetry]
 name = "tuhlbox"
-version = "0.4.0"
+version = "0.4.1"
 homepage = "https://git.uibk.ac.at/csak8736/tuhlbox"
 description = "Personal toolbox of language processing models."
 authors = ["Benjamin Murauer <b.murauer@posteo.de>"]

tests/test_torch_cnn.py

 """Test basic skorch models with CNN network."""
 
 import torch
-from dstoolbox.transformers import TextFeaturizer, Padder2d
+from dstoolbox.transformers import Padder2d, TextFeaturizer
 from sklearn.datasets import fetch_20newsgroups
 from sklearn.pipeline import make_pipeline
 from skorch import NeuralNetClassifier
-
+from tuhlbox.skorch_wrapper import TorchClassifier
 from tuhlbox.torch_cnn import CharCNN
 
 x, y = fetch_20newsgroups(return_X_y=True)
@@ -17,18 +17,19 @@ MAX_SEQ_LEN = 100
 pipe = make_pipeline(
     TextFeaturizer(max_features=VOCAB_SIZE),
     Padder2d(pad_value=VOCAB_SIZE, max_len=MAX_SEQ_LEN, dtype=int),
-    NeuralNetClassifier(
+    TorchClassifier(
         module=CharCNN,
-        device='cuda',
+        device="cuda",
         batch_size=54,
         max_epochs=5,
-        lr=0.01,
+        learn_rate=0.01,
         optimizer=torch.optim.Adam,
-        module__embedding_dim=EMB_DIM,
-        module__vocab_size=VOCAB_SIZE,
-        module__max_seq_length=MAX_SEQ_LEN,
-        module__num_classes=len(set(y)),
-    )
+        model_kwargs=dict(
+            module__embedding_dim=EMB_DIM,
+            module__vocab_size=VOCAB_SIZE,
+            module__max_seq_length=MAX_SEQ_LEN,
+        ),
+    ),
 )
 
 pipe.fit(x, y)

tests/test_torch_rnn.py

@@ -4,8 +4,7 @@ import torch
 from dstoolbox.transformers import Padder2d, TextFeaturizer
 from sklearn.datasets import fetch_20newsgroups
 from sklearn.pipeline import make_pipeline
-from skorch import NeuralNetClassifier
-from tuhlbox.torch_cnn import CharCNN
+from tuhlbox.skorch_wrapper import TorchClassifier
 from tuhlbox.torch_lstm import RNNClassifier
 
 x, y = fetch_20newsgroups(return_X_y=True)
@@ -17,12 +16,12 @@ MAX_SEQ_LEN = 100
 pipe = make_pipeline(
     TextFeaturizer(max_features=VOCAB_SIZE),
     Padder2d(pad_value=VOCAB_SIZE, max_len=MAX_SEQ_LEN, dtype=int),
-    NeuralNetClassifier(
+    TorchClassifier(
         module=RNNClassifier,
         device="cuda",
         batch_size=54,
         max_epochs=5,
-        lr=0.01,
+        learn_rate=0.01,
         optimizer=torch.optim.Adam,
     ),
 )

tuhlbox/skorch_wrapper.py

+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Optional, Type
+
+import torch
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.preprocessing import LabelEncoder
+from skorch import NeuralNetClassifier
+from torch import nn
+
+
+class TorchClassifier(ClassifierMixin, BaseEstimator):
+    """
+    Simple wrapper for Torch modules enabling string labels and dynamic classes.
+
+    Using this wrapper, the model does not need to know the number of classes that the
+    problem has. Instead, the number of classes is learned form the fit data.
+
+    Additionally, this wrapper is able to use string labels by using a LabelEncoder.
+    """
+
+    def __init__(
+        self,
+        module: Type[nn.Module],
+        batch_size: int = 64,
+        max_epochs: int = 5,
+        learn_rate: float = 1e-3,
+        device: str = "cuda",
+        model_kwargs: Dict[str, Any] = None,
+        optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
+    ):
+        self.module = module
+        self.device = device
+        self.batch_size = batch_size
+        self.max_epochs = max_epochs
+        self.learn_rate = learn_rate
+        self.model_kwargs = model_kwargs or {}
+        self.wrapped_model: Optional[NeuralNetClassifier] = None
+        self.optimizer = optimizer
+        self.label_encoder: LabelEncoder = LabelEncoder()
+
+    def fit(self, x: Any, y: Iterable[Any], **fit_kwargs: Any) -> TorchClassifier:
+        if self.wrapped_model is None:
+            classes = set(y)
+            n_classes = len(classes)
+            self.model_kwargs["module__n_classes"] = n_classes
+            self.wrapped_model = NeuralNetClassifier(
+                module=self.module,
+                device=self.device,
+                batch_size=self.batch_size,
+                max_epochs=self.max_epochs,
+                lr=self.learn_rate,
+                optimizer=self.optimizer,
+                classes=classes,
+                **self.model_kwargs
+            )
+            y = self.label_encoder.fit_transform(y)
+        self.wrapped_model.fit(x, y, **fit_kwargs)
+        return self
+
+    def predict(self, x: Any) -> Any:
+        if self.wrapped_model is None:
+            raise ValueError("model was not fitted")
+        predicted_classes = self.wrapped_model.predict(x)
+        return self.label_encoder.inverse_transform(predicted_classes)

tuhlbox/torch_cnn.py

@@ -10,7 +10,7 @@ class CharCNN(nn.Module):
         self,
         vocab_size: int,
         embedding_dim: int,
-        num_classes: int,
+        n_classes: int,
         max_seq_length: int,
         dropout: float = 0.0,
     ):
@@ -29,7 +29,7 @@ class CharCNN(nn.Module):
         self.vocab_size = vocab_size
         self.dropout = dropout
         self.max_seq_length = max_seq_length
-        self.num_classes = num_classes
+        self.num_classes = n_classes
 
         # Embedding Input dimensions (x, y):
         # x: batch size
@@ -54,8 +54,8 @@ class CharCNN(nn.Module):
             nn.Conv1d(
                 in_channels=self.embedding_dim,
                 out_channels=50,
-                kernel_size=7,
-                stride=1,
+                kernel_size=(7,),
+                stride=(1,),
             ),
             nn.ReLU(),
             nn.MaxPool1d(kernel_size=3, stride=3),
@@ -65,8 +65,8 @@ class CharCNN(nn.Module):
             nn.Conv1d(
                 in_channels=50,
                 out_channels=50,
-                kernel_size=5,
-                stride=1,
+                kernel_size=(5,),
+                stride=(1,),
             ),
             nn.ReLU(),
             nn.MaxPool1d(kernel_size=3, stride=3),

tuhlbox/torch_lstm.py

@@ -8,6 +8,7 @@ from torch import nn
 class RNNClassifier(nn.Module):
     def __init__(
         self,
+        n_classes: int,
         embedding_dim: int = 128,
         rec_layer_type: str = "lstm",
         num_units: int = 128,
@@ -21,6 +22,7 @@ class RNNClassifier(nn.Module):
         self.num_units = num_units
         self.num_layers = num_layers
         self.dropout = dropout
+        self.n_classes = n_classes
 
         self.emb = nn.Embedding(vocab_size + 1, embedding_dim=self.embedding_dim)
 
@@ -34,7 +36,7 @@ class RNNClassifier(nn.Module):
             batch_first=True,
         )
 
-        self.output = nn.Linear(self.num_units, 2)
+        self.output = nn.Linear(self.num_units, self.n_classes)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         embeddings = self.emb(x)