src/global_variables.py

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

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


def setup():
    global hf_api
    global all_metadata
    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)

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",
    )

if gr.NO_RELOAD:
    setup()