src/submission/submit.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# flake8: noqa E501

import json
import os
from datetime import datetime, timezone

from src.display.formatting import styled_error, styled_message, styled_warning
from src.envs import API, EVAL_REQUESTS_PATH, REQUESTS_REPO, TOKEN
from src.submission.check_validity import already_submitted_models, check_model_card, get_model_size, is_model_on_hub
from src.utils import get_request_hash, get_request_id
from src.submission.notify import notify

REQUESTED_MODELS = None
USERS_TO_SUBMISSION_DATES = None


def add_new_eval(
    model_name: str,
    # base_model: str,
    revision: str,
    # precision: str,
    # weight_type: str,
    model_type: str,
):
    # Remove if added to the form
    weight_type = 'Original'
    precision = 'Unknown'

    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if not REQUESTED_MODELS:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)

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

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

    # Is the model on the hub?
    # if weight_type in ["Delta", "Adapter"]:
    #     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 weight_type != "Adapter":
        model_on_hub, error, _ = is_model_on_hub(model_name=model_name, revision=revision, token=TOKEN, test_tokenizer=True)
        if not model_on_hub:
            return styled_error(f'Model "{model_name}" {error}')

    # Is the model info correctly filled?
    try:
        model_info = API.model_info(repo_id=model_name, revision=revision)
    except Exception:
        return styled_error("Could not get your model information. Please, fill it up properly.")

    model_size = get_model_size(
        model_info=model_info,
        precision=precision,
    )

    # Were the model card and license filled?
    try:
        license_title = model_info.cardData["license"]
    except Exception:
        license_title = "Unknown"
        # return styled_error("Please, select a license for your model.")

    is_model_card_ok, error_msg = check_model_card(model_name)
    if not is_model_card_ok:
        return styled_error(error_msg)

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

    eval_entry = {
        # "model": model,
        # "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "weight_type": weight_type,
        "status": "PENDING",
        "submitted_time": current_time,
        "model_type": model_type.split()[1],  # remove the emoji
        # "likes": model_info.likes,
        "params": model_size,
        "license": license_title,
        # "private": False,
    }

    # Check for duplicate submission
    request_id = get_request_id(model_name, revision, precision)
    if request_id in REQUESTED_MODELS:
        return styled_warning('This model has already been submitted.')
    request_hash = get_request_hash(model_name, revision, precision)

    print("Creating evaluation request file...")
    OUT_DIR = f"{EVAL_REQUESTS_PATH}/{model_name}"
    os.makedirs(OUT_DIR, exist_ok=True)

    out_path = f"{OUT_DIR}/{request_hash}.json"
    if os.path.exists(out_path):
        os.remove(out_path)

    with open(out_path, 'w') as f:
        json.dump(
            eval_entry,
            f,
            indent=4,
        )

    print("Uploading evaluation file...")
    API.upload_file(
        path_or_fileobj=out_path,
        path_in_repo='{}/{}.json'.format(model_name, request_hash),
        repo_id=REQUESTS_REPO,
        repo_type="dataset",
        commit_message=f"Add an evaluation request for {model_name}",
    )

    notify(model_name)

    # Remove the local file
    os.remove(out_path)

    return styled_message(
        "Your model has been submitted."
    )