registry refactor

llama_stack/providers/impls/meta_reference/evals/datas/dataset_registry.py

@@ -1,24 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-from .datasets import CustomDataset, HFDataset
-
-# TODO: make this into a config based registry
-DATASETS_REGISTRY = {
-    "mmlu-simple-eval-en": CustomDataset(
-        name="mmlu_eval",
-        url="https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv",
-    ),
-    "mmmu-accounting": HFDataset(
-        name="mmlu_eval",
-        url="hf://hellaswag?split=validation&trust_remote_code=True",
-    ),
-}
-
-
-def get_dataset(dataset_id: str):
-    dataset = DATASETS_REGISTRY[dataset_id]
-    dataset.load()
-    return dataset

llama_stack/providers/impls/meta_reference/evals/datas/datasets.py

@@ -1,62 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from abc import ABC, abstractmethod
-from urllib.parse import parse_qs, urlparse
-
-import pandas
-from datasets import Dataset, load_dataset
-
-
-class BaseDataset(ABC):
-    def __init__(self, name: str):
-        self.dataset = None
-        self.dataset_id = name
-        self.type = self.__class__.__name__
-
-    def __iter__(self):
-        return iter(self.dataset)
-
-    @abstractmethod
-    def load(self):
-        pass
-
-
-class CustomDataset(BaseDataset):
-    def __init__(self, name, url):
-        super().__init__(name)
-        self.url = url
-
-    def load(self):
-        if self.dataset:
-            return
-        # TODO: better support w/ data url
-        if self.url.endswith(".csv"):
-            df = pandas.read_csv(self.url)
-        elif self.url.endswith(".xlsx"):
-            df = pandas.read_excel(self.url)
-
-        self.dataset = Dataset.from_pandas(df)
-
-
-class HFDataset(BaseDataset):
-    def __init__(self, name, url):
-        super().__init__(name)
-        self.url = url
-
-    def load(self):
-        if self.dataset:
-            return
-
-        parsed = urlparse(self.url)
-
-        if parsed.scheme != "hf":
-            raise ValueError(f"Unknown HF dataset: {self.url}")
-
-        query = parse_qs(parsed.query)
-        query = {k: v[0] for k, v in query.items()}
-        path = parsed.netloc
-        self.dataset = load_dataset(path, **query)

llama_stack/providers/impls/meta_reference/evals/evals.py

@@ -8,11 +8,9 @@ from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.evals import *  # noqa: F403
 from termcolor import cprint
 
-from llama_stack.distribution.registry.tasks.task_registry import TaskRegistry
+from llama_stack.distribution.registry.datasets.dataset_registry import DatasetRegistry
 
-from llama_stack.providers.impls.meta_reference.evals.datas.dataset_registry import (
-    get_dataset,
-)
+from llama_stack.distribution.registry.tasks.task_registry import TaskRegistry
 
 from .config import MetaReferenceEvalsImplConfig
 
@@ -34,7 +32,8 @@ class MetaReferenceEvalsImpl(Evals):
         task: str,
     ) -> EvaluateResponse:
         cprint(f"model={model}, dataset={dataset}, task={task}", "red")
-        dataset = get_dataset(dataset)
+        dataset = DatasetRegistry.get_dataset(dataset)
+        dataset.load()
         task_impl = TaskRegistry.get_task(task)(dataset)
 
         x1 = task_impl.preprocess()