leaderboard/dataset.py

import os

import pandas as pd
from datasets import Dataset, load_dataset
from dotenv import load_dotenv
from huggingface_hub import login

load_dotenv()

HF_TOKEN = os.environ.get("HF_TOKEN", None)
RESULT_DATASET_ID = os.environ.get("RESULT_DATASET_ID", None)
if HF_TOKEN:
    login(token=HF_TOKEN)


def get_leaderboard_df() -> pd.DataFrame:
    # リーダーボードデータセットをHugging Face Hubからロード
    ds = load_dataset(RESULT_DATASET_ID, split="train")
    df = ds.to_pandas()
    return df


def save_leaderboard_df(df: pd.DataFrame):
    # DataFrameをDataset化
    ds = Dataset.from_pandas(df, preserve_index=False)
    # push_to_hubで更新
    ds.push_to_hub(RESULT_DATASET_ID, token=HF_TOKEN, commit_message="Update leaderboard")