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()