File size: 716 Bytes
f1f682e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from .base_dataset import BaseDataset
from .unobench import UnoBenchDataset

# Dataset name to class mapping
DATASET_REGISTRY = {
    "UNO-Bench": UnoBenchDataset,
    # Add more datasets here, for example:
    # "MMMU": MMMUDataset,
}

def get_dataset(dataset_name: str) -> BaseDataset:
    """
    Dataset factory function.

    :param dataset_name: Registered dataset name (e.g., 'UNO-Bench')
    :param split: Dataset split
    :return: Instance of BaseDataset
    """
    if dataset_name not in DATASET_REGISTRY:
        raise ValueError(f"Unknown dataset: {dataset_name}. Available benchmarks: {list(DATASET_REGISTRY.keys())}")
    
    dataset_class = DATASET_REGISTRY[dataset_name]
    return dataset_class()