src/huggingface/huggingface_client.py

from huggingface_hub import HfApi, login  # type: ignore[import]
import os
from loguru import logger
from .get_files import get_mp4_paths
from .get_environments import get_environments


class HuggingFaceClient:
    def __init__(self) -> None:
        login(  # type: ignore[no-untyped-call]
            token=os.environ.get("HUGGINGFACE_TOKEN"),
            add_to_git_credential=True,
            new_session=False,
        )
        self.hf_api = HfApi()
        self.environments = get_environments(self.hf_api)
        self.mp4_paths = get_mp4_paths(environments=self.environments)
        

    def api_predict(self, env_id: str)-> bytes|None:
        if env_id not in self.mp4_paths:
            logger.error(f"Environment {env_id} not found in {self.mp4_paths}")
            return None
        with open(self.mp4_paths[env_id], "rb") as f:
            return f.read()