Comrade3370/Speedway2

1import json2import os3from datetime import datetime, timezone4 5from src.display.formatting import styled_error, styled_message, styled_warning6from src.envs import API, EVAL_REQUESTS_PATH, TOKEN, QUEUE_REPO7from src.submission.check_validity import (8    already_submitted_models,9    check_model_card,10    get_model_size,11    is_model_on_hub,12)13 14REQUESTED_MODELS = None15USERS_TO_SUBMISSION_DATES = None16 17def add_new_eval(18    model: str,19    base_model: str,20    revision: str,21    precision: str,22    weight_type: str,23    model_type: str,24):25    global REQUESTED_MODELS26    global USERS_TO_SUBMISSION_DATES27    if not REQUESTED_MODELS:28        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)29 30    user_name = ""31    model_path = model32    if "/" in model:33        user_name = model.split("/")[0]34        model_path = model.split("/")[1]35 36    precision = precision.split(" ")[0]37    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")38 39    if model_type is None or model_type == "":40        return styled_error("Please select a model type.")41 42    # Does the model actually exist?43    if revision == "":44        revision = "main"45 46    # Is the model on the hub?47    if weight_type in ["Delta", "Adapter"]:48        base_model_on_hub, error, _ = is_model_on_hub(model_name=base_model, revision=revision, token=TOKEN, test_tokenizer=True)49        if not base_model_on_hub:50            return styled_error(f'Base model "{base_model}" {error}')51 52    if not weight_type == "Adapter":53        model_on_hub, error, _ = is_model_on_hub(model_name=model, revision=revision, token=TOKEN, test_tokenizer=True)54        if not model_on_hub:55            return styled_error(f'Model "{model}" {error}')56 57    # Is the model info correctly filled?58    try:59        model_info = API.model_info(repo_id=model, revision=revision)60    except Exception:61        return styled_error("Could not get your model information. Please fill it up properly.")62 63    model_size = get_model_size(model_info=model_info, precision=precision)64 65    # Were the model card and license filled?66    try:67        license = model_info.cardData["license"]68    except Exception:69        return styled_error("Please select a license for your model")70 71    modelcard_OK, error_msg = check_model_card(model)72    if not modelcard_OK:73        return styled_error(error_msg)74 75    # Seems good, creating the eval76    print("Adding new eval")77 78    eval_entry = {79        "model": model,80        "base_model": base_model,81        "revision": revision,82        "precision": precision,83        "weight_type": weight_type,84        "status": "PENDING",85        "submitted_time": current_time,86        "model_type": model_type,87        "likes": model_info.likes,88        "params": model_size,89        "license": license,90        "private": False,91    }92 93    # Check for duplicate submission94    if f"{model}_{revision}_{precision}" in REQUESTED_MODELS:95        return styled_warning("This model has been already submitted.")96 97    print("Creating eval file")98    OUT_DIR = f"{EVAL_REQUESTS_PATH}/{user_name}"99    os.makedirs(OUT_DIR, exist_ok=True)100    out_path = f"{OUT_DIR}/{model_path}_eval_request_False_{precision}_{weight_type}.json"101 102    with open(out_path, "w") as f:103        f.write(json.dumps(eval_entry))104 105    print("Uploading eval file")106    API.upload_file(107        path_or_fileobj=out_path,108        path_in_repo=out_path.split("eval-queue/")[1],109        repo_id=QUEUE_REPO,110        repo_type="dataset",111        commit_message=f"Add {model} to eval queue",112    )113 114    # Remove the local file115    os.remove(out_path)116 117    return styled_message(118        "Your request has been submitted to the evaluation queue!\nPlease wait for up to an hour for the model to show in the PENDING list."119    )120