model_interface/hf_model_store.py

import os
from pathlib import Path

from huggingface_hub import hf_hub_download
from huggingface_hub.errors import EntryNotFoundError

HF_REPO_ID = os.getenv("HF_REPO_ID", "Swaraj2003/test_model_2")
HF_REPO_TYPE = os.getenv("HF_REPO_TYPE", "model")
HF_REVISION = os.getenv("HF_REVISION")
HF_CACHE_DIR = os.getenv("HF_HOME")


def get_artifact_path(relative_path: str) -> Path:
    candidates = [relative_path]

    if relative_path.startswith("model_files/"):
        candidates.append(relative_path.removeprefix("model_files/"))
    else:
        candidates.append(f"model_files/{relative_path}")

    last_error = None
    for candidate in dict.fromkeys(candidates):
        try:
            return Path(
                hf_hub_download(
                    repo_id=HF_REPO_ID,
                    filename=candidate,
                    repo_type=HF_REPO_TYPE,
                    revision=HF_REVISION,
                    cache_dir=HF_CACHE_DIR,
                )
            )
        except EntryNotFoundError as exc:
            last_error = exc

    raise FileNotFoundError(
        f"Artifact not found in '{HF_REPO_ID}'. Tried paths: {candidates}. "
        "Verify repo folder structure and file names."
    ) from last_error