[Evals API][2/n] datasets / datasetio meta-reference implementation (#288)

* skeleton dataset / datasetio

* dataset datasetio

* config

* address comments

* delete dataset_utils

* address comments

* naming fix

llama_stack/providers/datatypes.py

@@ -10,6 +10,8 @@ from typing import Any, List, Optional, Protocol
 from llama_models.schema_utils import json_schema_type
 from pydantic import BaseModel, Field
 
+from llama_stack.apis.datasets import DatasetDef
+
 from llama_stack.apis.memory_banks import MemoryBankDef
 
 from llama_stack.apis.models import ModelDef
@@ -22,12 +24,14 @@ class Api(Enum):
     safety = "safety"
     agents = "agents"
     memory = "memory"
+    datasetio = "datasetio"
 
     telemetry = "telemetry"
 
     models = "models"
     shields = "shields"
     memory_banks = "memory_banks"
+    datasets = "datasets"
 
     # built-in API
     inspect = "inspect"
@@ -51,6 +55,12 @@ class MemoryBanksProtocolPrivate(Protocol):
     async def register_memory_bank(self, memory_bank: MemoryBankDef) -> None: ...
 
 
+class DatasetsProtocolPrivate(Protocol):
+    async def list_datasets(self) -> List[DatasetDef]: ...
+
+    async def register_datasets(self, dataset_def: DatasetDef) -> None: ...
+
+
 @json_schema_type
 class ProviderSpec(BaseModel):
     api: Api

llama_stack/providers/impls/meta_reference/datasetio/__init__.py

@@ -0,0 +1,18 @@
+# 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 .config import MetaReferenceDatasetIOConfig
+
+
+async def get_provider_impl(
+    config: MetaReferenceDatasetIOConfig,
+    _deps,
+):
+    from .datasetio import MetaReferenceDatasetIOImpl
+
+    impl = MetaReferenceDatasetIOImpl(config)
+    await impl.initialize()
+    return impl

llama_stack/providers/impls/meta_reference/datasetio/config.py

@@ -0,0 +1,9 @@
+# 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 llama_stack.apis.datasetio import *  # noqa: F401, F403
+
+
+class MetaReferenceDatasetIOConfig(BaseModel): ...

llama_stack/providers/impls/meta_reference/datasetio/datasetio.py

@@ -0,0 +1,142 @@
+# 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 typing import List, Optional
+
+import pandas
+
+from llama_models.llama3.api.datatypes import *  # noqa: F403
+
+from llama_stack.apis.datasetio import *  # noqa: F403
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from llama_stack.providers.datatypes import DatasetsProtocolPrivate
+
+from .config import MetaReferenceDatasetIOConfig
+
+
+class BaseDataset(ABC):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    @abstractmethod
+    def __len__(self) -> int:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def __getitem__(self, idx):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def load(self):
+        raise NotImplementedError()
+
+
+@dataclass
+class DatasetInfo:
+    dataset_def: DatasetDef
+    dataset_impl: BaseDataset
+
+
+class PandasDataframeDataset(BaseDataset):
+    def __init__(self, dataset_def: DatasetDef, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.dataset_def = dataset_def
+        self.df = None
+
+    def __len__(self) -> int:
+        assert self.df is not None, "Dataset not loaded. Please call .load() first"
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            return self.df.iloc[idx].to_dict(orient="records")
+        else:
+            return self.df.iloc[idx].to_dict()
+
+    def load(self) -> None:
+        if self.df is not None:
+            return
+
+        # TODO: more robust support w/ data url
+        if self.dataset_def.url.uri.endswith(".csv"):
+            df = pandas.read_csv(self.dataset_def.url.uri)
+        elif self.dataset_def.url.uri.endswith(".xlsx"):
+            df = pandas.read_excel(self.dataset_def.url.uri)
+        elif self.dataset_def.url.uri.startswith("data:"):
+            parts = parse_data_url(self.dataset_def.url.uri)
+            data = parts["data"]
+            if parts["is_base64"]:
+                data = base64.b64decode(data)
+            else:
+                data = unquote(data)
+                encoding = parts["encoding"] or "utf-8"
+                data = data.encode(encoding)
+
+            mime_type = parts["mimetype"]
+            mime_category = mime_type.split("/")[0]
+            data_bytes = io.BytesIO(data)
+
+            if mime_category == "text":
+                df = pandas.read_csv(data_bytes)
+            else:
+                df = pandas.read_excel(data_bytes)
+        else:
+            raise ValueError(f"Unsupported file type: {self.dataset_def.url}")
+
+        self.df = df
+
+
+class MetaReferenceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
+    def __init__(self, config: MetaReferenceDatasetIOConfig) -> None:
+        self.config = config
+        # local registry for keeping track of datasets within the provider
+        self.dataset_infos = {}
+
+    async def initialize(self) -> None: ...
+
+    async def shutdown(self) -> None: ...
+
+    async def register_dataset(
+        self,
+        dataset_def: DatasetDef,
+    ) -> None:
+        dataset_impl = PandasDataframeDataset(dataset_def)
+        self.dataset_infos[dataset_def.identifier] = DatasetInfo(
+            dataset_def=dataset_def,
+            dataset_impl=dataset_impl,
+        )
+
+    async def list_datasets(self) -> List[DatasetDef]:
+        return [i.dataset_def for i in self.dataset_infos.values()]
+
+    async def get_rows_paginated(
+        self,
+        dataset_id: str,
+        rows_in_page: int,
+        page_token: Optional[str] = None,
+        filter_condition: Optional[str] = None,
+    ) -> PaginatedRowsResult:
+        dataset_info = self.dataset_infos.get(dataset_id)
+        dataset_info.dataset_impl.load()
+
+        if page_token is None:
+            next_page_token = 0
+        else:
+            next_page_token = int(page_token)
+
+        if rows_in_page == -1:
+            rows = dataset_info.dataset_impl[next_page_token:]
+
+        start = next_page_token
+        end = min(start + rows_in_page, len(dataset_info.dataset_impl))
+        rows = dataset_info.dataset_impl[start:end]
+
+        return PaginatedRowsResult(
+            rows=rows,
+            total_count=len(rows),
+            next_page_token=str(end),
+        )

llama_stack/providers/registry/datasetio.py

@@ -0,0 +1,22 @@
+# 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 typing import List
+
+from llama_stack.distribution.datatypes import *  # noqa: F403
+
+
+def available_providers() -> List[ProviderSpec]:
+    return [
+        InlineProviderSpec(
+            api=Api.datasetio,
+            provider_type="meta-reference",
+            pip_packages=["pandas"],
+            module="llama_stack.providers.impls.meta_reference.datasetio",
+            config_class="llama_stack.providers.impls.meta_reference.datasetio.MetaReferenceDatasetIOConfig",
+            api_dependencies=[],
+        ),
+    ]

llama_stack/providers/tests/datasetio/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/tests/datasetio/provider_config_example.yaml

@@ -0,0 +1,4 @@
+providers:
+  - provider_id: test-meta
+    provider_type: meta-reference
+    config: {}

llama_stack/providers/tests/datasetio/test_datasetio.py

@@ -0,0 +1,109 @@
+# 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.
+import os
+
+import pytest
+import pytest_asyncio
+
+from llama_stack.apis.datasetio import *  # noqa: F403
+from llama_stack.distribution.datatypes import *  # noqa: F403
+from llama_stack.providers.tests.resolver import resolve_impls_for_test
+
+# How to run this test:
+#
+# 1. Ensure you have a conda with the right dependencies installed. This is a bit tricky
+#    since it depends on the provider you are testing. On top of that you need
+#    `pytest` and `pytest-asyncio` installed.
+#
+# 2. Copy and modify the provider_config_example.yaml depending on the provider you are testing.
+#
+# 3. Run:
+#
+# ```bash
+# PROVIDER_ID=<your_provider> \
+#   PROVIDER_CONFIG=provider_config.yaml \
+#   pytest -s llama_stack/providers/tests/datasetio/test_datasetio.py \
+#   --tb=short --disable-warnings
+# ```
+
+
+@pytest_asyncio.fixture(scope="session")
+async def datasetio_settings():
+    impls = await resolve_impls_for_test(
+        Api.datasetio,
+    )
+    return {
+        "datasetio_impl": impls[Api.datasetio],
+        "datasets_impl": impls[Api.datasets],
+    }
+
+
+async def register_dataset(datasets_impl: Datasets):
+    dataset = DatasetDefWithProvider(
+        identifier="test_dataset",
+        provider_id=os.environ["PROVIDER_ID"],
+        url=URL(
+            uri="https://openaipublic.blob.core.windows.net/simple-evals/mmlu.csv",
+        ),
+        columns_schema={},
+    )
+    await datasets_impl.register_dataset(dataset)
+
+
+@pytest.mark.asyncio
+async def test_datasets_list(datasetio_settings):
+    # NOTE: this needs you to ensure that you are starting from a clean state
+    # but so far we don't have an unregister API unfortunately, so be careful
+    datasets_impl = datasetio_settings["datasets_impl"]
+    response = await datasets_impl.list_datasets()
+    assert isinstance(response, list)
+    assert len(response) == 0
+
+
+@pytest.mark.asyncio
+async def test_datasets_register(datasetio_settings):
+    # NOTE: this needs you to ensure that you are starting from a clean state
+    # but so far we don't have an unregister API unfortunately, so be careful
+    datasets_impl = datasetio_settings["datasets_impl"]
+    await register_dataset(datasets_impl)
+
+    response = await datasets_impl.list_datasets()
+    assert isinstance(response, list)
+    assert len(response) == 1
+
+    # register same dataset with same id again will fail
+    await register_dataset(datasets_impl)
+    response = await datasets_impl.list_datasets()
+    assert isinstance(response, list)
+    assert len(response) == 1
+    assert response[0].identifier == "test_dataset"
+
+
+@pytest.mark.asyncio
+async def test_get_rows_paginated(datasetio_settings):
+    datasetio_impl = datasetio_settings["datasetio_impl"]
+    datasets_impl = datasetio_settings["datasets_impl"]
+    await register_dataset(datasets_impl)
+
+    response = await datasetio_impl.get_rows_paginated(
+        dataset_id="test_dataset",
+        rows_in_page=3,
+    )
+
+    assert isinstance(response.rows, list)
+    assert len(response.rows) == 3
+    assert response.next_page_token == "3"
+
+    # iterate over all rows
+    response = await datasetio_impl.get_rows_paginated(
+        dataset_id="test_dataset",
+        rows_in_page=10,
+        page_token=response.next_page_token,
+    )
+
+    assert isinstance(response.rows, list)
+    assert len(response.rows) == 10
+    assert response.next_page_token == "13"