src/global_variables.py

"""Global variables used in the space.
"""

import os
import json
from huggingface_hub import HfApi
import jsonlines

import gradio as gr

from src.constants import DATASET_NAME, HF_TOKEN, ASSETS_FOLDER

hf_api: HfApi
all_metadata: dict
all_votes: dict


def setup():
    global hf_api
    global all_metadata
    global all_votes
    hf_api = HfApi(token=HF_TOKEN)
    # hf_api.snapshot_download(
    #     local_dir=f"{ASSETS_FOLDER}/{DATASET_NAME}",
    #     repo_id=DATASET_NAME,
    #     repo_type="dataset",
    # )
    all_metadata = {}
    for split in ["train", "validation", "test"]:
        all_metadata[split] = []
        with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl") as reader:
            for row in reader:
                all_metadata[split].append(row)
    all_votes = {}
    for file in os.listdir(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes"):
        with open(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{file}") as f:
            key = file.split(".")[0]
            all_votes[key] = json.load(f)

def get_metadata(split):
    global all_metadata
    global hf_api
    hf_api.hf_hub_download(
        repo_id=DATASET_NAME,
        filename="metadata.jsonl",
        subfolder=f"data/{split}",
        repo_type="dataset",
        local_dir=f"{ASSETS_FOLDER}/{DATASET_NAME}",
    )
    all_metadata[split] = []
    with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl") as reader:
        for row in reader:
            all_metadata[split].append(row)

def save_metadata(split):
    global all_metadata
    values = all_metadata[split]
    with jsonlines.open(f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl", mode='w') as writer:
        writer.write_all(values)
    hf_api.upload_file(
        path_or_fileobj=f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/metadata.jsonl",
        path_in_repo=f"data/{split}/metadata.jsonl",
        repo_id=DATASET_NAME,
        repo_type="dataset",
    )

def get_votes(key):
    global all_votes
    global hf_api
    try:
        hf_api.hf_hub_download(
            repo_id=DATASET_NAME,
            filename=f"votes/{key}.json",
            repo_type="dataset",
            local_dir=f"{ASSETS_FOLDER}/{DATASET_NAME}",
        )
        with open(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{key}.json") as f:
            all_votes[key] = json.load(f)
    except:
        pass

def save_votes(key):
    global all_votes
    global hf_api
    with open(f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{key}.json", "w") as f:
        json.dump(all_votes[key], f)
    hf_api.upload_file(
        path_or_fileobj=f"{ASSETS_FOLDER}/{DATASET_NAME}/votes/{key}.json",
        path_in_repo=f"votes/{key}.json",
        repo_id=DATASET_NAME,
        repo_type="dataset",
    )

if gr.NO_RELOAD:
    setup()