Hf model name download (#1048)

* first commit

* add test

* typo

* typo

* ruff

* add hf readonly fg token for tests

* address review

* raise error if type not recognized

* add pydantic to reqs

* typo

* correctly revert

---------

Co-authored-by: lbluque <[email protected]>

Author: zulissimeta

.github/workflows/test.yml

@@ -53,6 +53,8 @@ jobs:
           echo "$(readlink -f .)" >> $GITHUB_PATH
 
       - name: Test core with pytest
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN_OMAT_READONLY }}
         run: |
           pytest tests -vv --ignore=tests/demo/ocpapi/tests/integration/  --cov-report=xml --cov=fairchem -c ./packages/fairchem-core/pyproject.toml
 

packages/env.cpu.yml

@@ -26,4 +26,5 @@ dependencies:
 - submitit
 - tensorboard
 - wandb
+- huggingface_hub
 name: fair-chem

packages/env.gpu.yml

@@ -29,4 +29,5 @@ dependencies:
 - submitit
 - tensorboard
 - wandb
+- huggingface_hub
 name: fair-chem

packages/fairchem-core/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "fairchem-core"
-description = "Machine learning models for use in catalysis as part of the Open Catalyst Project"
+description = "Machine learning models for chemistry and materials science by the FAIR Chemistry team"
 license = {text = "MIT License"}
 dynamic = ["version", "readme"]
 requires-python = ">=3.9, <3.13"
@@ -24,7 +24,9 @@ dependencies = [
     "tqdm",
     "submitit",
     "hydra-core",
-    "torchtnt"
+    "torchtnt",
+    "huggingface_hub>=0.29.2",
+    "pydantic>=2.10.0"
 ]
 
 [project.optional-dependencies]  # add optional dependencies to be installed as pip install fairchem.core[dev]

packages/requirements.txt

@@ -1,3 +1,4 @@
 torch==2.4.1
 numpy==1.26.4
 ase==3.24.0
+huggingface_hub==0.29.2

src/fairchem/core/models/model_registry.py

@@ -7,26 +7,43 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import os
 import shutil
 from importlib import resources
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import requests
-import yaml
+from huggingface_hub import hf_hub_download
+from pydantic import AnyUrl, BaseModel
 
 from fairchem.core import models
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
-with (resources.files(models) / "pretrained_models.yml").open("rt") as f:
-    MODEL_REGISTRY = yaml.safe_load(f)
+class HuggingFaceModel(BaseModel):
+    type: Literal["huggingface_hub"]
+    repo_id: Literal["fairchem/OMAT24"]
+    filename: str
 
 
-available_pretrained_models = tuple(MODEL_REGISTRY.keys())
+class URLModel(BaseModel):
+    url: str
+    type: Literal["url"]
+
+
+class ModelRegistry(BaseModel):
+    models: dict[str, AnyUrl | HuggingFaceModel | URLModel]
+
+
+with (resources.files(models) / "pretrained_models.json").open("rb") as f:
+    MODEL_REGISTRY = ModelRegistry(models=json.load(f))
+
+
+available_pretrained_models = tuple(MODEL_REGISTRY.models.keys())
 
 
 def model_name_to_local_file(model_name: str, local_cache: str | Path) -> str:
@@ -40,25 +57,40 @@ def model_name_to_local_file(model_name: str, local_cache: str | Path) -> str:
         str: local path to checkpoint file
     """
     logging.info(f"Checking local cache: {local_cache} for model {model_name}")
-    if model_name not in MODEL_REGISTRY:
+    if model_name not in available_pretrained_models:
         logging.error(f"Not a valid model name '{model_name}'")
         raise ValueError(
             f"Not a valid model name '{model_name}'. Model name must be one of {available_pretrained_models}"
         )
-    if not os.path.exists(local_cache):
-        os.makedirs(local_cache, exist_ok=True)
-    if not os.path.exists(local_cache):
-        logging.error(f"Failed to create local cache folder '{local_cache}'")
-        raise RuntimeError(f"Failed to create local cache folder '{local_cache}'")
-    model_url = MODEL_REGISTRY[model_name]
-    local_path = os.path.join(local_cache, os.path.basename(model_url))
-
-    # download the file
-    if not os.path.isfile(local_path):
-        local_path_tmp = local_path + ".tmp"  # download to a tmp file in case we fail
-        with open(local_path_tmp, "wb") as out:
-            response = requests.get(model_url, stream=True)
-            response.raw.decode_content = True
-            shutil.copyfileobj(response.raw, out)
-        shutil.move(local_path_tmp, local_path)
-    return local_path
+
+    if isinstance(MODEL_REGISTRY.models[model_name], URLModel):
+        # We have a url to download
+
+        if not os.path.exists(local_cache):
+            os.makedirs(local_cache, exist_ok=True)
+        if not os.path.exists(local_cache):
+            logging.error(f"Failed to create local cache folder '{local_cache}'")
+            raise RuntimeError(f"Failed to create local cache folder '{local_cache}'")
+        model_url = MODEL_REGISTRY.models[model_name].url
+        local_path = os.path.join(local_cache, os.path.basename(model_url))
+
+        # download the file
+        if not os.path.isfile(local_path):
+            local_path_tmp = (
+                local_path + ".tmp"
+            )  # download to a tmp file in case we fail
+            with open(local_path_tmp, "wb") as out:
+                response = requests.get(model_url, stream=True)
+                response.raw.decode_content = True
+                shutil.copyfileobj(response.raw, out)
+            shutil.move(local_path_tmp, local_path)
+        return local_path
+    elif isinstance(MODEL_REGISTRY.models[model_name], HuggingFaceModel):
+        return hf_hub_download(
+            repo_id=MODEL_REGISTRY.models[model_name].repo_id,
+            filename=MODEL_REGISTRY.models[model_name].filename,
+        )
+    else:
+        raise NotImplementedError(
+            f"{type(MODEL_REGISTRY.models[model_name])} is an unknown registry type."
+        )