base torch

Author: SeBorgey

autointent/_dump_tools/main.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import numpy.typing as npt
+import torch
 
 from autointent.configs import CrossEncoderConfig, EmbedderConfig
 from autointent.context.optimization_info import Artifact
@@ -108,6 +109,8 @@ def dump(
                 simple_attrs[key] = val
             elif isinstance(val, np.ndarray):
                 arrays[key] = val
+            elif isinstance(val, torch.Tensor):
+                arrays[key] = val.cpu().numpy()
             else:
                 # Use the appropriate dumper for complex objects
                 Dumper._dump_single_object(key, val, path, exists_ok, raise_errors)

autointent/modules/scoring/_gcn/gcn_model.py

@@ -4,8 +4,10 @@
 import torch
 import torch.nn as nn
 from pydantic import BaseModel
+from typing_extensions import Self
 
 from autointent._utils import detect_device
+from autointent._wrappers import BaseTorchModuleWithVocab
 
 
 class GCNModelDumpMetadata(BaseModel):
@@ -28,7 +30,7 @@ def forward(self, adj_matrix, features):
         return output
 
 
-class TextMLGCN(nn.Module):
+class TextMLGCN(BaseTorchModuleWithVocab):
     _metadata_dict_name = "metadata.json"
     _state_dict_name = "state_dict.pt"
 
@@ -41,7 +43,7 @@ def __init__(
         p_reweight: float,
         tau_threshold: float,
     ):
-        super().__init__()
+        super().__init__(embed_dim=bert_feature_dim)
         self.num_classes = num_classes
         self.p_reweight = p_reweight
         self.tau_threshold = tau_threshold
@@ -93,7 +95,7 @@ def set_correlation_matrix(self, train_labels):
         )
         self.correlation_matrix.data.copy_(corr_matrix)
 
-    def forward(self, bert_features, label_embeddings):
+    def forward(self, bert_features, label_embeddings):  # type: ignore
         classifiers = label_embeddings
         for i in range(len(self.gcn_layers)):
             classifiers = self.gcn_layers[i](self.correlation_matrix, classifiers)
@@ -102,10 +104,6 @@ def forward(self, bert_features, label_embeddings):
         logits = torch.matmul(bert_features, classifiers.T)
         return logits
 
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
     def dump(self, path: Path) -> None:
         metadata = GCNModelDumpMetadata(
             num_classes=self.num_classes,
@@ -124,7 +122,7 @@ def dump(self, path: Path) -> None:
         self.to(device)
 
     @classmethod
-    def load(cls, path: Path, device: str | None = None) -> "TextMLGCN":
+    def load(cls, path: Path, device: str | None = None) -> Self:
         with (path / cls._metadata_dict_name).open() as file:
             metadata = GCNModelDumpMetadata(**json.load(file))
         device = device or detect_device()

autointent/modules/scoring/_gcn/gcn_scorer.py

@@ -6,9 +6,10 @@
 from pydantic import PositiveInt
 from torch import nn
 from torch.utils.data import DataLoader, TensorDataset
+from typing_extensions import Self
 
 from autointent import Context, Embedder
-from autointent.configs import EmbedderConfig, TaskTypeEnum, TorchTrainingConfig
+from autointent.configs import CrossEncoderConfig, EmbedderConfig, TaskTypeEnum, TorchTrainingConfig
 from autointent.custom_types import ListOfLabels
 from autointent.modules.base import BaseScorer
 from autointent.modules.scoring._gcn.gcn_model import TextMLGCN
@@ -162,3 +163,15 @@ def clear_cache(self) -> None:
         if hasattr(self, "_label_embedder"):
             self._label_embedder.clear_ram()
             del self._label_embedder
+
+    @classmethod
+    def load(
+        cls,
+        path: str,
+        embedder_config: EmbedderConfig | None = None,
+        cross_encoder_config: CrossEncoderConfig | None = None,
+    ) -> Self:
+        instance = super().load(path, embedder_config, cross_encoder_config)
+        if hasattr(instance, "_label_embeddings"):
+            instance._label_embeddings = torch.tensor(instance._label_embeddings).to(instance.torch_config.device)
+        return instance