fix datasetio

This commit is contained in:
Xi Yan 2024-11-11 15:29:29 -05:00
parent acd055d763
commit e27c6e3662
5 changed files with 15 additions and 15 deletions

View file

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

View file

@ -6,4 +6,4 @@
from llama_stack.apis.datasetio import * # noqa: F401, F403
class MetaReferenceDatasetIOConfig(BaseModel): ...
class LocalfsDatasetIOConfig(BaseModel): ...

View file

@ -15,7 +15,7 @@ from dataclasses import dataclass
from llama_stack.providers.datatypes import DatasetsProtocolPrivate
from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
from .config import MetaReferenceDatasetIOConfig
from .config import LocalfsDatasetIOConfig
class BaseDataset(ABC):
@ -77,8 +77,8 @@ class PandasDataframeDataset(BaseDataset):
self.df = self._validate_dataset_schema(df)
class MetaReferenceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
def __init__(self, config: MetaReferenceDatasetIOConfig) -> None:
class LocalfsDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
def __init__(self, config: LocalfsDatasetIOConfig) -> None:
self.config = config
# local registry for keeping track of datasets within the provider
self.dataset_infos = {}

View file

@ -13,10 +13,10 @@ def available_providers() -> List[ProviderSpec]:
return [
InlineProviderSpec(
api=Api.datasetio,
provider_type="meta-reference",
provider_type="localfs",
pip_packages=["pandas"],
module="llama_stack.providers.inline.datasetio.meta_reference",
config_class="llama_stack.providers.inline.datasetio.meta_reference.MetaReferenceDatasetIOConfig",
module="llama_stack.providers.inline.datasetio.localfs",
config_class="llama_stack.providers.inline.datasetio.localfs.LocalfsDatasetIOConfig",
api_dependencies=[],
),
remote_provider_spec(

View file

@ -19,12 +19,12 @@ def datasetio_remote() -> ProviderFixture:
@pytest.fixture(scope="session")
def datasetio_meta_reference() -> ProviderFixture:
def datasetio_localfs() -> ProviderFixture:
return ProviderFixture(
providers=[
Provider(
provider_id="meta-reference",
provider_type="meta-reference",
provider_id="localfs",
provider_type="localfs",
config={},
)
],
@ -44,7 +44,7 @@ def datasetio_huggingface() -> ProviderFixture:
)
DATASETIO_FIXTURES = ["meta_reference", "remote", "huggingface"]
DATASETIO_FIXTURES = ["localfs", "remote", "huggingface"]
@pytest_asyncio.fixture(scope="session")