test hf hub

mlip_arena/models/__init__.py

@@ -27,19 +27,42 @@ class MLIP(
         super().__init__(*args, **kwargs)
 
 
+class ModuleMLIP(MLIP):
+    def __init__(self, model: nn.Module, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.register_module("model", model)
+
+    def forward(self, x):
+        print("Forwarding...")
+        out = self.model(x)
+        print("Forwarded!")
+        return out
+
+
 class MLIPCalculator(Calculator):
+    name: str
+    device: torch.device
+    model: MLIP
+    implemented_properties: list[str] = ["energy", "forces", "stress"]
+
     def __init__(
         self,
+        # PyTorch
         model_path: str | Path,
         device: torch.device | None = None,
+        # ASE Calculator
+        restart=None,
+        atoms=None,
+        directory=".",
+        **kwargs,
     ):
-        super().__init__()
-        self.name: str = self.__class__.__name__
-        self.device = device or torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
-        self.model: MLIP = MLIP.from_pretrained(model_path, map_location=self.device)
-        self.implemented_properties = ["energy", "forces", "stress"]
+        super().__init__(restart=restart, atoms=atoms, directory=directory, **kwargs)
+        # self.name: str = self.__class__.__name__
+        # self.device = device or torch.device(
+        #     "cuda" if torch.cuda.is_available() else "cpu"
+        # )
+        # self.model: MLIP = MLIP.from_pretrained(model_path, map_location=self.device)
+        # self.implemented_properties = ["energy", "forces", "stress"]
 
     def calculate(
         self, atoms: Atoms, properties: list[str], system_changes: list = all_changes

mlip_arena/models/mace.py

@@ -4,20 +4,54 @@ from ase.calculators.calculator import all_changes
 from huggingface_hub import hf_hub_download
 from torch_geometric.data import Data
 
-from mlip_arena.models import MLIPCalculator
+from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP
 
 
 class MACE_MP_Medium(MLIPCalculator):
-    def __init__(self, device: torch.device | None = None):
+    def __init__(
+        self,
+        device: torch.device | None = None,
+        restart=None,
+        atoms=None,
+        directory=".",
+        **kwargs,
+    ):
+        # Download the pytorch model from huggingface to local and load it
+        # NOTE: this is not the ideal way to load the model, but it is the simplest
+        # way to do it for now. Ideally, if the model is the subclass of PyTorchModelHubMixin,
+        # we should be able to load it directly from the hub or local using MLIP class.
         fpath = hf_hub_download(
             repo_id="cyrusyc/mace-universal",
             subfolder="pretrained",
             filename="2023-12-12-mace-128-L1_epoch-199.model",
             revision=None,  # TODO: Add revision
         )
-        super().__init__(model_path=fpath, device=device)
+        # module = ModuleMLIP(torch.load(fpath, map_location="cpu"))
+        print(torch.load(fpath, map_location="cpu"))
+        repo_id = f"atomind/{self.__class__.__name__}".replace("_", "-")
+        # module.save_pretrained(
+        #     save_directory=self.__class__.__name__,
+        #     repo_id=repo_id,
+        #     push_to_hub=True,
+        # )
+
+        super().__init__(
+            model_path=repo_id,
+            device=device,
+            restart=restart,
+            atoms=atoms,
+            directory=directory,
+            **kwargs,
+        )
+
+        # self.name: str = self.__class__.__name__
+        # self.device = device or torch.device(
+        #     "cuda" if torch.cuda.is_available() else "cpu"
+        # )
+        # self.model: MLIP = ModuleMLIP.from_pretrained(repo_id, map_location=self.device)
+        # self.implemented_properties = ["energy", "forces", "stress"]
 
-        self.name = "MACE-MP-0 (medium)"
+        self.display = "MACE-MP-0 (medium)"
         self.version = "1.0.0"
         self.implemented_properties = [
             "energy",

pyproject.toml

@@ -31,6 +31,7 @@ dependencies=[
     "torch_dftd>=0.4.0",
     "huggingface_hub",
     "torch-geometric>=2.5.2",
+    "safetensors"
 ]
 
 [project.urls]

tests/hf_hub.ipynb

@@ -0,0 +1,100 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from huggingface_hub import hf_hub_download\n",
+    "from mlip_arena.models import MLIP, MLIPCalculator, ModuleMLIP"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "fpath = hf_hub_download(\n",
+    "    repo_id=\"cyrusyc/mace-universal\",\n",
+    "    subfolder=\"pretrained\",\n",
+    "    filename=\"2023-12-12-mace-128-L1_epoch-199.model\",\n",
+    "    revision=None,  # TODO: Add revision\n",
+    ")\n",
+    "\n",
+    "model = torch.load(fpath, map_location=\"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "module = ModuleMLIP(model=model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "model.safetensors: 100%|██████████| 44.2M/44.2M [00:02<00:00, 20.6MB/s]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "CommitInfo(commit_url='https://huggingface.co/atomind/mace-mp-medium/commit/ef94f6bd9c7167bb28d594d5a3e7a5fbfbda2acb', commit_message='Push model using huggingface_hub.', commit_description='', oid='ef94f6bd9c7167bb28d594d5a3e7a5fbfbda2acb', pr_url=None, pr_revision=None, pr_num=None)"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "module.save_pretrained(\n",
+    "    \"mace\",\n",
+    "    repo_id=\"atomind/MACE_MP_Medium\".lower().replace(\"_\", \"-\"),\n",
+    "    push_to_hub=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}