meg/backend

1import glob2import json3from dataclasses import dataclass4# from datetime import datetime, timezone5from typing import Optional6 7from huggingface_hub import HfApi, snapshot_download8 9from src.envs import TOKEN10from src.logging import setup_logger11 12logger = setup_logger(__name__)13 14PENDING_STATUS = "PENDING"15RUNNING_STATUS = "RUNNING"16FINISHED_STATUS = "FINISHED"17FAILED_STATUS = "FAILED"18 19 20@dataclass21class EvalRequest:22    """This class represents one evaluation request file.23    """24    model: str25    status: str26    json_filepath: str27    weight_type: str = "Original"28    precision: str = ""  # float16, bfloat1629    revision: str = "main"  # commit hash30    submitted_time: Optional[31        str] = "2022-05-18T11:40:22.519222"  # random date just so that we can still order requests by date32    model_type: Optional[str] = None  # pretrained, fine-tuned, etc33    likes: Optional[int] = 034    params: Optional[int] = None35    license: Optional[str] = ""36    base_model: Optional[str] = ""37    private: Optional[bool] = False38 39    def get_model_args(self):40        """Edit this function if you want to manage more complex quantization issues. You'll need to map it to 41        the evaluation suite you chose.42        """43        model_args = f"pretrained={self.model},revision={self.revision}"44 45        if self.precision in ["float16", "bfloat16", "float32"]:46            model_args += f",dtype={self.precision}"47 48        # Quantized models need some added config, the install of bits and bytes, etc49 50        # elif self.precision == "8bit":51        #    model_args += ",load_in_8bit=True"52        # elif self.precision == "4bit":53        #    model_args += ",load_in_4bit=True"54        # elif self.precision == "GPTQ":55        # A GPTQ model does not need dtype to be specified,56        # it will be inferred from the config57        else:58            raise Exception(f"Unknown precision {self.precision}.")59 60        return model_args61 62 63def set_eval_request(api: HfApi, eval_request: EvalRequest, set_to_status: str,64                     hf_repo: str, local_dir: str):65    """Updates a given eval request with its new status on the hub (running, completed, failed, ...)"""66    json_filepath = eval_request.json_filepath67 68    with open(json_filepath) as fp:69        data = json.load(fp)70 71    data["status"] = set_to_status72 73    with open(json_filepath, "w") as f:74        f.write(json.dumps(data))75 76    api.upload_file(77        path_or_fileobj=json_filepath,78        path_in_repo=json_filepath.replace(local_dir, ""),79        repo_id=hf_repo,80        repo_type="dataset",81    )82 83 84def get_eval_requests(local_dir: str, hf_repo: str) -> list[EvalRequest]:85    """Gets all pending evaluation requests and return a list in which private86    models appearing first, followed by public models sorted by the number of87    likes.88 89    Returns:90        `list[EvalRequest]`: a list of model info dicts.91    """92    snapshot_download(repo_id=hf_repo, revision="main", local_dir=local_dir,93                      repo_type="dataset", max_workers=60, token=TOKEN)94    json_files = glob.glob(f"{local_dir}/**/*.json", recursive=True)95 96    eval_requests = []97    for json_filepath in json_files:98        with open(json_filepath) as fp:99            data = json.load(fp)100        if data["status"] in [PENDING_STATUS, RUNNING_STATUS]:101            data["json_filepath"] = json_filepath102            eval_request = EvalRequest(**data)103            eval_requests.append(eval_request)104 105    return eval_requests106 107 108def check_set_to_fail(eval_request: EvalRequest):109    """Checks whether a file says it's RUNNING to determine whether to FAIL"""110    json_filepath = eval_request.json_filepath111 112    with open(json_filepath) as fp:113        data = json.load(fp)114 115    status = data["status"]116    # Don't fail pending tasks.117    if status == PENDING_STATUS:118        return False119    else:120        return True121        # time_format = "%Y-%m-%dT%H:%M:%SZ"122        # submitted_time_str = data["submitted_time"]123        # submitted_time_naive = datetime.strptime(submitted_time_str,124        #                                         time_format)125        # current_time = datetime.now(126        #    timezone.utc)  # .strftime("%Y-%m-%dT%H:%M:%SZ")127        # submitted_time = submitted_time_naive.replace(128        #    tzinfo=current_time.tzinfo)129        # difference = current_time - submitted_time130        # diff_seconds = difference.total_seconds()131        # If it's been running for less than 2 hours, leave it alone.132        # if diff_seconds < 7200:133        #    return False134        # else:135        #    return True136 137 138def check_completed_evals(139        api: HfApi,140        hf_repo: str,141        local_dir: str,142        completed_status: str,143        failed_status: str,144        hf_repo_results: str,145        local_dir_results: str,146):147    """Checks if the currently running evals are completed, if yes, update their status on the hub."""148    snapshot_download(repo_id=hf_repo_results, revision="main",149                      local_dir=local_dir_results, repo_type="dataset",150                      max_workers=60, token=TOKEN)151 152    eval_requests = get_eval_requests(hf_repo=hf_repo, local_dir=local_dir)153 154    for eval_request in eval_requests:155        model = eval_request.model156        logger.info("====================================")157        logger.info(f"Checking {model}")158 159        output_path = model160        output_file = f"{local_dir_results}/{output_path}/results*.json"161        output_file_exists = len(glob.glob(output_file)) > 0162 163        if output_file_exists:164            logger.info(165                f"EXISTS output file exists for {model} setting it to {completed_status}"166            )167            set_eval_request(api, eval_request, completed_status, hf_repo,168                             local_dir)169        else:170            set_to_fail = check_set_to_fail(eval_request)171            if set_to_fail:172                logger.info(173                    f"No result file found for {model} setting it to {failed_status}"174                )175                set_eval_request(api, eval_request, failed_status, hf_repo,176                                 local_dir)177