File size: 1,398 Bytes
5b344d4
 
5ae8333
7011484
 
 
76c534f
5ae8333
 
7011484
5ae8333
 
 
 
 
 
 
 
 
7011484
 
 
 
 
 
51cf3cb
5ae8333
 
7011484
 
 
 
 
 
51cf3cb
5ae8333
 
 
7011484
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from __future__ import annotations

from pathlib import Path

import yaml
from huggingface_hub import hf_hub_download

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 Path(local_env_results_path).open() as f:
        env_results = yaml.safe_load(f)
    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 Path(local_model_results_path).open() as f:
        model_results = yaml.safe_load(f)

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