Create model.py

model.py

@@ -0,0 +1,105 @@
+import json
+import re
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from jaxtyping import Float
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class Model(nn.Module, PyTorchModelHubMixin, repo_url="https://huggingface.co/TimWalter/RAM/new/main", license="mit",):
+    """
+    RAM model that predicts reachability from a modified Denavit-Hartenberg morphology parametrisation and a pose.
+    """
+
+    @classmethod
+    def from_id(cls, model_id: int):
+        """
+        Instantiate a model from its wandb ID.
+
+        Args:
+            model_id: Wandb ID of the model.
+
+        Returns:
+            model: Instantiated model.
+        """
+
+        model_dir = Path(__file__).parent.parent / "trained_models"
+        pattern = rf"{model_id}-[a-z]+-[a-z]+"
+        folder = next((f for f in model_dir.iterdir() if re.match(pattern, f.name)), None)
+        metadata_path = model_dir / folder / 'metadata.json'
+        metadata = json.load(open(metadata_path, 'r'))
+
+        model = cls(**metadata["hyperparameter"])
+        model_folder = Path(str(model_dir / folder))
+        prime = model_folder / "model.pth"
+
+        if prime.exists():
+            old_state_dict =torch.load(prime)
+        else:
+            old_state_dict =torch.load(model_folder / "checkpoint.pth")
+        new_state_dict = {}
+        # 2. Map legacy state dict keys to your new flattened structure
+        for key, value in old_state_dict.items():
+            new_key = key
+            if key.startswith("encoder.lstm."):
+                new_key = key.replace("encoder.lstm.", "encoder.")
+            elif key.startswith("decoder.model."):
+                new_key = key.replace("decoder.model.", "decoder.")
+
+            new_state_dict[new_key] = value
+
+        # Load the corrected state dict into the model
+        model.load_state_dict(new_state_dict)
+
+        return model
+
+    def __init__(self,
+                 dim_encoding: int = 128,
+                 num_encoder_layers: int = 1,
+                 drop_prob: float = 0.0,
+                 dim_decoder: int = 1792,
+                 num_decoder_layer: int = 8):
+        """
+        Initialise the model.
+
+        Args:
+            dim_encoding: The dimension of the latent morphology encoding.
+            num_encoder_layers: Number of LSTM layers.
+            drop_prob: Dropout probability of the LSTM.
+            dim_decoder: Hidden dimension of the MLP.
+            num_decoder_layer: Number of layers of the MLP.
+        """
+
+        super().__init__()
+        self.encoder = nn.LSTM(3, dim_encoding, num_encoder_layers, dropout=drop_prob, batch_first=True, bias=False)
+        self.decoder = nn.Sequential(
+            nn.Linear(9 + dim_encoding, dim_decoder),
+            nn.ReLU(),
+            *[nn.Sequential(nn.Linear(dim_decoder, dim_decoder), nn.ReLU())
+              for _ in range(num_decoder_layer)],
+            nn.Linear(dim_decoder, 1)
+        )
+
+    def forward(self, morph: Float[Tensor, "batch seq 3"], pose: Float[Tensor, "batch 9"]) -> Float[Tensor, "batch"]:
+        """
+        Predict reachability.
+
+        Args:
+            morph: Morphology description.
+            pose: Pose as vector encoded.
+        Returns:
+            Reachability logit.
+        """
+        latent = self.encoder(morph)[1][0][-1]
+        logit = self.decoder(torch.cat([pose, latent], dim=-1)).squeeze(-1)
+        return logit
+
+    @torch.inference_mode()
+    def predict(self, *args, **kwargs):
+        """
+        forward with inference mode
+        """
+        return self.forward(*args, **kwargs)