utils.py

import pandas as pd


LEADERBOARD_PATH = f"chan0369/muse_leaderboard"

DEFAULT_METRICS = [
    'VerbMem Forget',
    'KnowMem Forget',
    'PrivLeak',
    'KnowMem Retain'
]

CORPORA = ['bbc', 'hp']


def load_data(corpus):
    assert corpus in CORPORA
    df = pd.read_csv(f"submission/{corpus}.csv")
    df = df[['Method', 'Submitted By'] + DEFAULT_METRICS + ['Id']]
    return df