cleanup hardcoded dataset registry

This commit is contained in:
Xi Yan 2024-10-14 14:19:15 -07:00
parent a9210cd416
commit 9c501d042b
3 changed files with 31 additions and 29 deletions

View file

@ -3,32 +3,9 @@
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
# TODO: make these import config based
from llama_stack.apis.datasets import * # noqa: F403
from ..registry import Registry
from .dataset_wrappers import CustomDataset, HuggingfaceDataset
class DatasetRegistry(Registry[BaseDataset]):
_REGISTRY: Dict[str, BaseDataset] = {}
DATASETS_REGISTRY = [
CustomDataset(
config=CustomDatasetDef(
identifier="mmlu-simple-eval-en",
url="https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv",
)
),
HuggingfaceDataset(
config=HuggingfaceDatasetDef(
identifier="hellaswag",
dataset_name="hellaswag",
kwargs={"split": "validation", "trust_remote_code": True},
)
),
]
for d in DATASETS_REGISTRY:
DatasetRegistry.register(d.dataset_id, d)