Upload folder using huggingface_hub

src/constants.py

@@ -0,0 +1,6 @@
+SPACE_REPO = "c-gohlke/litrl"
+SPACE_REPO_TYPE = "space"
+MODEL_REPO = "c-gohlke/litrl"
+MODEL_REPO_TYPE = "model"
+
+ENV_RESULTS_FILE_DEPTH = 3

src/huggingface/get_environments.py

@@ -1,6 +1,7 @@
+from huggingface_hub import HfApi  # type: ignore[import]
+from src.constants import MODEL_REPO, MODEL_REPO_TYPE, ENV_RESULTS_FILE_DEPTH
 
-
-def get_environments() -> list[str]:
+def get_environments(hf_api: HfApi) -> list[str]:
     environments: list[str] = []
     files = hf_api.list_repo_files(MODEL_REPO, repo_type=MODEL_REPO_TYPE)
     for file in files:

src/huggingface/get_files.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+from huggingface_hub import hf_hub_download  # type: ignore[import]
+import yaml  # type: ignore[import]
+from src.constants import MODEL_REPO, MODEL_REPO_TYPE
+
+def get_mp4_paths(environments: list[str]) -> dict[str, Path]:
+    mp4_paths: dict[str, Path] = {}
+    for env in environments:
+        mp4_paths[env] = get_best(env, filename="demo.mp4")
+    return mp4_paths
+
+
+def get_best(env: str, filename: str = "demo.mp4") -> Path:
+    hf_env_results_path = f"models/{env}/results.yaml"
+    local_env_results_path = hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=hf_env_results_path)
+    with open(local_env_results_path, "r") as f:
+        env_results = yaml.load(f, Loader=yaml.FullLoader)
+    best_model_type = max(env_results, key=lambda model: env_results[model])
+    model_results_path = f"models/{env}/{best_model_type}/results.yaml"
+    local_model_results_path = hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=model_results_path)
+    with open(local_model_results_path, "r") as f:
+        model_results = yaml.load(f, Loader=yaml.FullLoader)
+
+    best_model = max(model_results, key=lambda model: model_results[model])
+    hf_gif_path = f"models/{env}/{best_model_type}/{best_model}/{filename}"
+    return Path(hf_hub_download(repo_id=MODEL_REPO, repo_type=MODEL_REPO_TYPE, filename=hf_gif_path))

src/huggingface/huggingface_client.py

@@ -1,18 +1,8 @@
-from huggingface_hub import hf_hub_download
-from huggingface_hub import HfApi, login
+from huggingface_hub import HfApi, login  # type: ignore[import]
 import os
-import yaml
-from pathlib import Path
 from loguru import logger
-from PIL import Image
-from .get_best_mp4 import get_mp4_paths
-
-SPACE_REPO = "c-gohlke/litrl"
-SPACE_REPO_TYPE = "space"
-MODEL_REPO = "c-gohlke/litrl"
-MODEL_REPO_TYPE = "model"
-
-ENV_RESULTS_FILE_DEPTH = 3
+from .get_files import get_mp4_paths
+from .get_environments import get_environments
 
 
 class HugingFaceClient:
@@ -23,7 +13,8 @@ class HugingFaceClient:
             new_session=False,
         )
         self.hf_api = HfApi()
-        self.mp4_paths = get_mp4_paths()
+        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:

src/typing.py

@@ -17,8 +17,4 @@ class RolloutPolicy(enum.Enum):
 class CpuConfig(BaseModel):
     agent_type: AgentType
     simulations: int | None = None
-    rollout_policy: RolloutPolicy | None = None
-
-    # def __format__(self, __format_spec: str) -> str:
-    #     raise ValueError(f"__format__ not implemented for {self.__class__.__name__}")
-    #     return super().__format__(__format_spec)
+    rollout_policy: RolloutPolicy | None = None