src/utils.py

from dataclasses import dataclass
import pandas as pd


@dataclass
class ColumnContent:
    name: str
    type: str
    displayed_by_default: bool
    hidden: bool = False


def fields(raw_class):
    return [v for k, v in raw_class.__dict__.items() if k[:2] != "__" and k[-2:] != "__"]


@dataclass(frozen=True)
class AutoEvalColumn:  # Auto evals column
    model = ColumnContent("Models", "markdown", True)
    re2text_easy = ColumnContent("Re2Text-Easy", "number", True)
    text2re_easy = ColumnContent("Text2Re-Easy", "number", True)
    re2text_hard = ColumnContent("Re2Text-Hard", "number", True)
    text2re_hard = ColumnContent("Text2Re-Hard", "number", True)
    avg = ColumnContent("Avg", "number", True)
    model_size = ColumnContent("Model Size", "markdown", True)

    link = ColumnContent("Links", "str", False)


def model_hyperlink(link, model_name):
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{model_name}</a>'


def make_clickable_names(df):
    df["Models"] = df.apply(
        lambda row: model_hyperlink(row["Links"], row["Models"]), axis=1
    )
    return df


def make_plot_data(df, task):
    c = []
    x = []
    y = []

    for i in df.index:
        c.append(df.loc[i, "pure_name"])
        x.append(f"{task}-Easy")
        y.append(df.loc[i, f"{task}-Easy"])

        c.append(df.loc[i, "pure_name"])
        x.append(f"{task}-Hard")
        y.append(df.loc[i, f"{task}-Hard"])

    data = pd.DataFrame(
        {
            "Symbol": c,
            "Setting": x,
            "Accuracy": y,
        }
    )

    return data