EmbodiedAgentInterface/backend

1import glob2import json3from dataclasses import dataclass4from typing import Optional5 6from huggingface_hub import HfApi, snapshot_download7from src.envs import TOKEN8from src.logging import setup_logger9 10logger = setup_logger(__name__)11 12PENDING_STATUS = "PENDING"13RUNNING_STATUS = "RUNNING"14FINISHED_STATUS = "FINISHED"15FAILED_STATUS = "FAILED"16 17@dataclass18class EvalRequest:19    """This class represents one evaluation request file.20    """21    model: str22    status: str23    json_filepath: str24    weight_type: str = "Original"25    model_type: str = ""  # pretrained, finetuned, with RL26    precision: str = ""  # float16, bfloat1627    revision: str = "main" # commit hash28    submitted_time: Optional[str] = "2022-05-18T11:40:22.519222"  # random date just so that we can still order requests by date29    model_type: Optional[str] = None # pretrained, fine-tuned, etc - define your own categories in 30    likes: Optional[int] = 031    params: Optional[int] = None32    license: Optional[str] = ""33    base_model: Optional[str] = ""34    private: Optional[bool] = False35    36    def get_model_args(self):37        """Edit this function if you want to manage more complex quantization issues. You'll need to map it to 38        the evaluation suite you chose.39        """40        model_args = f"pretrained={self.model},revision={self.revision}"41 42        if self.precision in ["float16", "bfloat16"]:43            model_args += f",dtype={self.precision}"44 45        # Quantized models need some added config, the install of bits and bytes, etc46        else:47            raise Exception(f"Unknown precision {self.precision}.")48        49        return model_args50 51 52def set_eval_request(api: HfApi, eval_request: EvalRequest, set_to_status: str, hf_repo: str, local_dir: str):53    """Updates a given eval request with its new status on the hub (running, completed, failed, ...)"""54    json_filepath = eval_request.json_filepath55 56    with open(json_filepath) as fp:57        data = json.load(fp)58 59    data["status"] = set_to_status60 61    with open(json_filepath, "w") as f:62        f.write(json.dumps(data))63 64    api.upload_file(65        path_or_fileobj=json_filepath,66        path_in_repo=json_filepath.replace(local_dir, ""),67        repo_id=hf_repo,68        repo_type="dataset",69    )70 71 72def get_eval_requests(job_status: list, local_dir: str, hf_repo: str) -> list[EvalRequest]:73    """Gets all pending evaluation requests and return a list in which private74    models appearing first, followed by public models sorted by the number of75    likes.76 77    Returns:78        `list[EvalRequest]`: a list of model info dicts.79    """80    snapshot_download(repo_id=hf_repo, revision="main", local_dir=local_dir, repo_type="dataset", max_workers=60, token=TOKEN)81    json_files = glob.glob(f"{local_dir}/**/*.json", recursive=True)82 83    eval_requests = []84    for json_filepath in json_files:85        with open(json_filepath) as fp:86            data = json.load(fp)87        if data["status"] in job_status:88            data["json_filepath"] = json_filepath89            eval_request = EvalRequest(**data)90            eval_requests.append(eval_request)91 92    return eval_requests93 94 95def eval_was_running(eval_request: EvalRequest):96    """Checks whether a file says it's RUNNING to determine whether to FAIL"""97    json_filepath = eval_request.json_filepath98 99    with open(json_filepath) as fp:100        data = json.load(fp)101 102    status = data["status"]103    return status == RUNNING_STATUS104 105def check_completed_evals(106    api: HfApi,107    hf_repo: str,108    local_dir: str,109    checked_status: str,110    completed_status: str,111    failed_status: str,112    hf_repo_results: str,113    local_dir_results: str,114):115    """Checks if the currently running evals are completed, if yes, update their status on the hub."""116    snapshot_download(117        repo_id=hf_repo_results, 118        revision="main", 119        local_dir=local_dir_results, 120        repo_type="dataset", 121        max_workers=60, 122        token=TOKEN123    )124 125    running_evals = get_eval_requests(checked_status, hf_repo=hf_repo, local_dir=local_dir)126 127    for eval_request in running_evals:128        model = eval_request.model129        logger.info("====================================")130        logger.info(f"Checking {model}")131 132        output_path = model133        output_file = f"{local_dir_results}/{output_path}/results*.json"134        output_file_exists = len(glob.glob(output_file)) > 0135 136        if output_file_exists:137            logger.info(138                f"EXISTS output file exists for {model} setting it to {completed_status}"139            )140            set_eval_request(api, eval_request, completed_status, hf_repo, local_dir)141        else:142            if eval_was_running(eval_request=eval_request):143                logger.info(144                    f"No result file found for {model} setting it to {failed_status}"145                )146                set_eval_request(api, eval_request, failed_status, hf_repo, local_dir)147