data/dataset_handler.py

VIDORE_DATASETS_KEYWORDS = [
    "arxivqa",
    "docvqa",
    "infovqa",
    "tabfquad",
    "tatdqa",
    "shift",
    "artificial_intelligence",
    "energy",
    "government_reports",
    "healthcare_industry",
]


def get_datasets_nickname(dataset_name) -> str:
    if "arxivqa" in dataset_name:
        return "ArxivQA"

    elif "docvqa" in dataset_name:
        return "DocVQA"
    elif "infovqa" in dataset_name:
        return "InfoVQA"
    elif "tabfquad" in dataset_name:
        return "TabFQuad"

    elif "tatdqa" in dataset_name:
        return "TAT-DQA"

    elif "shift" in dataset_name:
        return "Shift Project"

    elif "artificial_intelligence" in dataset_name:
        return "Artificial Intelligence"

    elif "energy" in dataset_name:
        return "Energy"

    elif "government_reports" in dataset_name:
        return "Government Reports"

    elif "healthcare_industry" in dataset_name:
        return "Healthcare Industry"

    else:
        raise ValueError(f"Dataset {dataset_name} not found in ViDoRe")