utils/prechecks.py

"""

"""

from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download

def check_hf_persistent_storage(
    repo_id: str = None,
    repo_type: str = "model",
    file_or_folder="file",
    target: str = None,
    destination: str = "/data/"
):

    file_path = Path(destination) / target

    def _download_file():
        try:
            if file_or_folder == "file":
                hf_hub_download(
                    repo_id=repo_id,
                    repo_type=repo_type,
                    filename=target,
                    local_dir=destination
                )
            elif file_or_folder == "folder":
                snapshot_download(
                    repo_id=repo_id,
                    repo_type=repo_type,
                    allow_patterns=f"{target}/**",
                    local_dir=destination
                )

            print(f"Successfully downloaded '{target}' to '{destination}'.")
        except Exception as e:
            print(f"An error occurred during the download: {e}")

    # Check if the file exists at the specified path
    if not file_path.exists():
        _download_file()
    else:
        print(f"File '{file_path}' already exists. No download needed.")