utils.py

import re
import json
import time
from huggingface_hub import HfApi


def current_seconds_time():
    return round(time.time())


def form_file_name(model_name, commit_id, inference_function):
    return f"predictions_{re.sub('/', '_', model_name)}_{commit_id}_{inference_function}.json"


def update_model_queue(repo_id, model_name, commit_id, inference_function, status):
    assert status in ["queued", "in_progress", "failed (online)"]
    api = HfApi()

    timestamp = current_seconds_time()
    predictions_filename = form_file_name(model_name, commit_id, inference_function)

    predictions_object = {
        "model_name": model_name,
        "predictions": [[""]],
        "commit_id": commit_id,
        "inference_function": inference_function,
        "last_updated_timestamp": timestamp,
        "status": status,
    }

    with open(predictions_filename, "w") as f:
        json.dump(predictions_object, f)

    future = api.upload_file(
        path_or_fileobj=predictions_filename,
        path_in_repo=predictions_filename,
        repo_id=repo_id,
        repo_type="dataset",
        run_as_future=True,
    )


def upload_predictions(repo_id, predictions, model_name, commit_id, inference_function):
    api = HfApi()

    timestamp = current_seconds_time()
    predictions_filename = form_file_name(model_name, commit_id, inference_function)

    predictions_object = {
        "model_name": model_name,
        "predictions": predictions,
        "commit_id": commit_id,
        "inference_function": inference_function,
        "last_updated_timestamp": timestamp,
        "status": "completed",
    }

    with open(predictions_filename, "w") as f:
        json.dump(predictions_object, f)

    future = api.upload_file(
        path_or_fileobj=predictions_filename,
        path_in_repo=predictions_filename,
        repo_id=repo_id,
        repo_type="dataset",
        run_as_future=True,
    )