added dump-load

Author: riapush

autointent/_dump_tools.py

@@ -102,12 +102,13 @@ def dump(obj: Any, path: Path) -> None:  # noqa: ANN401, C901, PLR0912, PLR0915
                 model_path = path / Dumper.torch_models / key
                 model_path.mkdir(parents=True, exist_ok=True)
                 try:
-                    torch.save(val.state_dict(), model_path / "model.pt")
-                    # Save class info for loading
+                    torch.save(val._model.state_dict(), model_path / "model.pt")
+                    vocab_path = path / Dumper.torch_models / "vocab.json"
+                    with vocab_path.open("w") as f:
+                        json.dump(obj._vocab, f)
                     class_info = {
                         "module": val.__class__.__module__,
                         "name": val.__class__.__name__,
-                        "is_textcnn": isinstance(val, TextCNN)
                     }
                     with (model_path / "class_info.json").open("w") as f:
                         json.dump(class_info, f)
@@ -254,25 +255,15 @@ def load(  # noqa: PLR0912, C901, PLR0915
                     try:
                         with (model_dir / "class_info.json").open("r") as f:
                             class_info = json.load(f)
+                        vocab_path = path / Dumper.torch_models / "vocab.json"
+                        with vocab_path.open("r") as f:
+                            obj._vocab = json.load(f)
 
                         module = __import__(class_info["module"], fromlist=[class_info["name"]])
                         model_class = getattr(module, class_info["name"])
 
                         # Create model instance
-                        if class_info.get("is_textcnn"):
-                            # For TextCNN, we need to get the parameters from the parent CNNScorer
-                            model = model_class(
-                                vocab_size=len(obj._vocab) if hasattr(obj, "_vocab") and obj._vocab else 0,
-                                n_classes=obj._n_classes if hasattr(obj, "_n_classes") else 0,
-                                embed_dim=obj.embed_dim if hasattr(obj, "embed_dim") else 128,
-                                kernel_sizes=obj.kernel_sizes if hasattr(obj, "kernel_sizes") else [3, 4, 5],
-                                num_filters=obj.num_filters if hasattr(obj, "num_filters") else 100,
-                                dropout=obj.dropout if hasattr(obj, "dropout") else 0.1,
-                                padding_idx=obj._pad_idx if hasattr(obj, "_pad_idx") else 0
-                            )
-                        else:
-                            # For other torch models, create with default parameters
-                            model = model_class()
+                        model = model_class()
 
                         # Load state dict
                         model.load_state_dict(torch.load(model_dir / "model.pt"))

autointent/modules/scoring/_cnn/textcnn.py

@@ -10,8 +10,8 @@ class TextCNN(nn.Module):
 
     def __init__(
         self,
-        vocab_size: int,
-        n_classes: int,
+        vocab_size: int = 0,
+        n_classes: int = 0,
         embed_dim: int = 128,
         kernel_sizes: list[int] = [3, 4, 5], # noqa: B006
         num_filters: int = 100,
@@ -22,15 +22,28 @@ def __init__(
         """Initialize TextCNN model."""
         super().__init__()
 
+        # Register model hyperparameters as buffers
+        self.register_buffer("vocab_size", torch.tensor(vocab_size))
+        self.register_buffer("n_classes", torch.tensor(n_classes))
+        self.register_buffer("embed_dim", torch.tensor(embed_dim))
+        self.register_buffer("kernel_sizes", torch.tensor(kernel_sizes))
+        self.register_buffer("num_filters", torch.tensor(num_filters))
+        self.register_buffer("dropout_rate", torch.tensor(dropout))
+        self.register_buffer("padding_idx", torch.tensor(padding_idx))
+
         if pretrained_embs is not None:
             _, embed_dim = pretrained_embs.shape
-            self.embedding = nn.Embedding.from_pretrained(pretrained_embs, freeze=True) # type: ignore[no-untyped-call]
+            self.embedding = nn.Embedding.from_pretrained(pretrained_embs, freeze=True)  # type: ignore[no-untyped-call]
+            # Register pretrained embeddings as buffer if they exist
+            self.register_buffer("pretrained_embs", pretrained_embs)
         else:
             self.embedding = nn.Embedding(
                 num_embeddings=vocab_size,
                 embedding_dim=embed_dim,
                 padding_idx=padding_idx
             )
+            # Register None for pretrained_embs buffer
+            self.register_buffer("pretrained_embs", None)
 
         self.convs = nn.ModuleList([
             nn.Conv1d(
@@ -50,3 +63,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         concatenated: torch.Tensor = torch.cat(conved, dim=1)
         dropped: torch.Tensor = self.dropout(concatenated)
         return self.fc(dropped) # type: ignore[no-any-return]
+    
+    def load(self, model_path: str) -> None:
+        """Load model from saved state.
+        
+        Args:
+            model_path: Path to the saved model state dictionary.
+        """
+        state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+        self.load_state_dict(state_dict)