utils.py

import pandas as pd

LEADERBOARD_PATH = f"chan0369/muse_leaderboard"

DEFAULT_METRICS = [
    'verbmem_f',
    'knowmem_f',
    'privleak',
    'knowmem_r',
]

DEFAULT_COLUMNS = ['name', 'organization'] + DEFAULT_METRICS + ['id']

DEFAULT_COLUMN_LABELS = [
    'Method Name',
    'Submitted By',
    'VerbMem ⬇️',
    'KnowMem Forget ⬇️',
    'PrivLeak',
    'KnowMem Retain (Utility) ⬆',
    'Submission Id'
]

CORPORA = ['news', 'books']


def load_data(corpus):
    assert corpus in CORPORA
    df = pd.read_csv(f"submission/{corpus}.csv")
    df = df[DEFAULT_COLUMNS].rename(columns={
        k: v for k, v in zip(DEFAULT_COLUMNS, DEFAULT_COLUMN_LABELS)
    })
    return df