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 # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
from .config import MetaReferenceDatasetIOConfig from .config import LocalfsDatasetIOConfig
async def get_provider_impl( async def get_provider_impl(
config: MetaReferenceDatasetIOConfig, config: LocalfsDatasetIOConfig,
_deps, _deps,
): ):
from .datasetio import MetaReferenceDatasetIOImpl from .datasetio import LocalfsDatasetIOImpl
impl = MetaReferenceDatasetIOImpl(config) impl = LocalfsDatasetIOImpl(config)
await impl.initialize() await impl.initialize()
return impl return impl

View file

@ -6,4 +6,4 @@
from llama_stack.apis.datasetio import * # noqa: F401, F403 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.datatypes import DatasetsProtocolPrivate
from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
from .config import MetaReferenceDatasetIOConfig from .config import LocalfsDatasetIOConfig
class BaseDataset(ABC): class BaseDataset(ABC):
@ -77,8 +77,8 @@ class PandasDataframeDataset(BaseDataset):
self.df = self._validate_dataset_schema(df) self.df = self._validate_dataset_schema(df)
class MetaReferenceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate): class LocalfsDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
def __init__(self, config: MetaReferenceDatasetIOConfig) -> None: def __init__(self, config: LocalfsDatasetIOConfig) -> None:
self.config = config self.config = config
# local registry for keeping track of datasets within the provider # local registry for keeping track of datasets within the provider
self.dataset_infos = {} self.dataset_infos = {}

View file

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

View file

@ -19,12 +19,12 @@ def datasetio_remote() -> ProviderFixture:
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def datasetio_meta_reference() -> ProviderFixture: def datasetio_localfs() -> ProviderFixture:
return ProviderFixture( return ProviderFixture(
providers=[ providers=[
Provider( Provider(
provider_id="meta-reference", provider_id="localfs",
provider_type="meta-reference", provider_type="localfs",
config={}, 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") @pytest_asyncio.fixture(scope="session")