Upload folder using huggingface_hub

equiformer_ase_calculator.py

@@ -42,7 +42,7 @@ class EquiformerASECalculator(Calculator):
 
     def __init__(
         self,
-        checkpoint_path: str,
+        checkpoint_path: Optional[str] = None,
         device: Optional[torch.device] = None,
         project_root: str = None,
         **kwargs,
@@ -69,6 +69,9 @@ class EquiformerASECalculator(Calculator):
             config = yaml.safe_load(file)
         model_config = config["model"]
         self.model = EquiformerV2_OC20(**model_config)
+
+        if checkpoint_path is None:
+            checkpoint_path = os.path.join(project_root, "ckpt/eqv2.ckpt")
         state_dict = torch.load(checkpoint_path, weights_only=True)["state_dict"]
         state_dict = {k.replace("potential.", ""): v for k, v in state_dict.items()}
         self.model.load_state_dict(state_dict, strict=False)

ocpmodels/checkpointing.py

@@ -2,12 +2,13 @@ import torch
 import itertools
 import warnings
 from collections import OrderedDict
-from typing import Any, List, Mapping, namedtuple   
+from typing import Any, List, Mapping, namedtuple
 
 # from torch/nn/modules/module.py
 
 _EXTRA_STATE_KEY_SUFFIX = "_extra_state"
 
+
 class _IncompatibleKeys(
     namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]),
 ):
@@ -18,6 +19,7 @@ class _IncompatibleKeys(
 
     __str__ = __repr__
 
+
 def _load_from_state_dict(
     self,
     state_dict,
@@ -199,8 +201,10 @@ def _load_from_state_dict(
                 elif input_name[0] not in local_state:
                     unexpected_keys.append(key)
 
+
 def load_state_dict(
-    self, state_dict: Mapping[str, Any], 
+    self,
+    state_dict: Mapping[str, Any],
     assign: bool = False,
     strict_unexpected_keys: bool = False,
     strict_missing_keys: bool = False,
@@ -242,9 +246,7 @@ def load_state_dict(
         ``RuntimeError``.
     """
     if not isinstance(state_dict, Mapping):
-        raise TypeError(
-            f"Expected state_dict to be dict-like, got {type(state_dict)}."
-        )
+        raise TypeError(f"Expected state_dict to be dict-like, got {type(state_dict)}.")
 
     missing_keys: List[str] = []
     unexpected_keys: List[str] = []
@@ -307,7 +309,7 @@ def load_state_dict(
                     ", ".join(f'"{k}"' for k in unexpected_keys[:20])
                 ),
             )
-    
+
     if strict_missing_keys:
         if len(missing_keys) > 0:
             error_msgs.insert(
@@ -323,4 +325,4 @@ def load_state_dict(
                 self.__class__.__name__, "\n\t".join(error_msgs)
             )
         )
-    return _IncompatibleKeys(missing_keys, unexpected_keys)
+    return _IncompatibleKeys(missing_keys, unexpected_keys)