utils.py

import pathlib
import tempfile
import json

import gradio as gr
import pandas as pd
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from constants import API, SUBMISSIONS_REPO, RESULTS_REPO, ASSAY_RENAME, LEADERBOARD_RESULTS_COLUMNS

pd.set_option('display.max_columns', None)

# def make_user_clickable(name):
#     link =f'https://huggingface.co/{name}'
#     return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{name}</a>'

# def make_boundary_clickable(filename):
#     link =f'https://huggingface.co/datasets/proxima-fusion/constellaration-bench-results/blob/main/{filename}'
#     return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">link</a>'


def show_output_box(message):
    return gr.update(value=message, visible=True)


def fetch_hf_results():
    # Should cache by default if not using force_redownload
    df = load_dataset(
        RESULTS_REPO, data_files="auto_submissions/metrics_all.csv",
    )["train"].to_pandas()
    assert all(col in df.columns for col in LEADERBOARD_RESULTS_COLUMNS), f"Expected columns {LEADERBOARD_RESULTS_COLUMNS} not found in {df.columns}. Missing columns: {set(LEADERBOARD_COLUMNS) - set(df.columns)}"
    df = df.drop_duplicates(subset=["model", "assay"])
    df["property"] = df["assay"].map(ASSAY_RENAME)
    print(df.head())
    return df


def read_result_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=RESULTS_REPO,
        repo_type="dataset",
        filename=filename,
    )
    return local_path


def read_submission_from_hub(filename):
    local_path = hf_hub_download(
        repo_id=SUBMISSIONS_REPO,
        repo_type="dataset",
        filename=filename,
    )
    return local_path


def write_results(record, result):
    record.update(result)
    record["result_filename"] = (
        record["submission_filename"].rstrip(".json") + "_results.json"
    )
    print(record["result_filename"])
    record["evaluated"] = True

    record["objectives"] = json.dumps(record.get("objectives", []))
    record["feasibilities"] = json.dumps(record.get("feasibility", []))

    if "objective" not in record.keys():
        record["objective"] = 0.0
        record["minimize_objective"] = True
        record["feasibility"] = sum(record["feasibility"]) / len(record["feasibility"])

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp:
        json.dump(record, tmp, indent=2)
        tmp.flush()
        tmp_name = tmp.name

    API.upload_file(
        path_or_fileobj=tmp_name,
        path_in_repo=record["result_filename"],
        repo_id=RESULTS_REPO,
        repo_type="dataset",
        commit_message=f"Add result data for {record['result_filename']}",
    )

    pathlib.Path(tmp_name).unlink()
    return


def get_user(profile: gr.OAuthProfile | None) -> str:
    if profile is None:
        return "Please login to submit a boundary for evaluation."
    return profile.username