src/submission/submit.py

import json
import os
import ast
from datetime import datetime, timezone

from src.display.formatting import styled_error, styled_message, styled_warning
from src.envs import API, EVAL_REQUESTS_PATH, TOKEN, QUEUE_REPO, PRIVATE_REPO
from src.submission.check_validity import (
    already_submitted_models,
    check_model_card,
    get_model_size,
    is_model_on_hub,
)
from src.display.utils import PromptTemplateName

REQUESTED_MODELS = None
USERS_TO_SUBMISSION_DATES = None

PLACEHOLDER_DATASET_WISE_NORMALIZATION_CONFIG = """{
    "NCBI" : {
        "" : "condition"
    },
    "CHIA" : {
        "" : "condition"
        "" : "drug"
        "" : "procedure"
        "" : "measurement"
    },
    "BIORED" : {
        "" : "condition"
        "" : "drug"
        "" : "gene"
        "" : "gene variant"
    },
    "BC5CDR" : {
        "" : "condition"
        "" : "drug"
    }
}

"""

def add_new_eval(
    model: str,
    base_model: str,
    revision: str,
    model_type: str,
    domain_specific: bool,
    chat_template: bool,
    precision: str,
    weight_type: str,
):
    """
    Saves request if valid else returns the error.
    Validity is checked based on -
        - model's existence on hub
        - necessary info on the model's card
        - label normalization is a valid python dict and contains the keys for all datasets
        - threshold for gliner is a valid float

    """
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if not REQUESTED_MODELS and not PRIVATE_REPO:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)
    if model.startswith("/"):
        if not PRIVATE_REPO:
            return styled_error("Private models are not allowed to be submitted to the public queue.")
        user_name = ""
        model_path = model
        private = True
    else:
        user_name = ""
        model_path = model
        if "/" in model:
            user_name = model.split("/")[0]
            model_path = model.split("/")[1]
        private = False

    # precision = precision.split(" ")[0]
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")

    if model_type is None or model_type == "":
        return styled_error("Please select a model type.")
    
    model_type = model_type.split(":")[-1].strip()

    # Does the model actually exist?
    if revision == "":
        revision = "main"

    # Is the model on the hub?
    if weight_type in ["Delta", "Adapter"] and not private:
        base_model_on_hub, error, _ = is_model_on_hub(model_name=base_model, revision=revision, token=TOKEN, test_tokenizer=True)
        if not base_model_on_hub:
            return styled_error(f'Base model "{base_model}" {error}')

    if not weight_type == "Adapter" and not private:
        model_on_hub, error, _ = is_model_on_hub(model_name=model, revision=revision, token=TOKEN, test_tokenizer=True)
        if not model_on_hub:
            return styled_error(f'Model "{model}" {error}')
    

    # Is the model info correctly filled?
    try:
        if not private:
            model_info = API.model_info(repo_id=model, revision=revision)
            model_size = get_model_size(model_info=model_info)
            license = model_info.cardData["license"]
            modelcard_OK, error_msg = check_model_card(model)
            if not modelcard_OK:
                return styled_error(error_msg)
            likes = model_info.likes
        else:
            model_size = None
            license = None
            likes = -1
    except Exception:
        return styled_error("Could not get your model information. Please fill it up properly.")

    
    # Verify the inference config now
    # try:
    #     label_normalization_map = ast.literal_eval(label_normalization_map)
    # except Exception as e:
    #     return styled_error("Please enter a valid json for the labe; normalization map")

    # inference_config = {
    #     # "model_arch" : model_arch,
    #     "label_normalization_map": label_normalization_map,
    # }

    # Seems good, creating the eval
    print("Adding new eval")


    eval_entry = {
        "model_name": model,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "weight_type": weight_type,
        "is_domain_specific": domain_specific,
        "use_chat_template": chat_template,
        "status": {
            "closed-ended": "PENDING",
            "open-ended": "PENDING",
            "med-safety": "PENDING",
            "medical-summarization": "PENDING",
            "note-generation": "PENDING",
        },
        "submitted_time": current_time,
        "model_type": model_type,
        "likes": likes,
        "num_params": model_size,
        "license": license,
        "private": private,
        "slurm_id": None
    }

    # Check for duplicate submission

    if not PRIVATE_REPO and f"{model}_{revision}_{precision}" in REQUESTED_MODELS:
        return styled_warning("This model has been already submitted. Add the revision if the model has been updated.")

    print("Creating eval file")
    if not private:
        OUT_DIR = f"{EVAL_REQUESTS_PATH}/{user_name}"
        
    else:
        OUT_DIR = f"{EVAL_REQUESTS_PATH}/offline"
        model_path = model_path[1:] if model_path.startswith("/") else model_path
        model_path = model_path.replace("/", "+-+")
        out_path = f"{OUT_DIR}/{model_path}_{revision}_{precision}_{weight_type}_eval_request.json"
    os.makedirs(OUT_DIR, exist_ok=True)
    with open(out_path, "w") as f:
        f.write(json.dumps(eval_entry))

    print("Uploading eval file")
    API.upload_file(
        path_or_fileobj=out_path,
        path_in_repo=out_path.split(f"{EVAL_REQUESTS_PATH}/")[1],
        repo_id=QUEUE_REPO,
        repo_type="dataset",
        commit_message=f"Add {model} to eval queue",
    )

    # Remove the local file
    os.remove(out_path)

    return styled_message(
        "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."
    )