phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-08 11:07:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

registry refactor

Browse source
This commit is contained in:
Xi Yan 2024-10-14 00:12:46 -07:00
parent 78cb88c3c4
commit 18fe966e96
9 changed files with 60 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

7
llama_stack/distribution/registry/datasets/__init__.py
View file

@ -6,8 +6,13 @@
# TODO: make these import config based
from llama_stack.apis.dataset import * # noqa: F403
from ..registry import Registry
from .dataset import CustomDataset, HuggingfaceDataset
from .dataset_registry import DatasetRegistry
class DatasetRegistry(Registry[BaseDataset]):
_REGISTRY: Dict[str, BaseDataset] = {}
DATASETS_REGISTRY = [
CustomDataset(

32
llama_stack/distribution/registry/datasets/dataset_registry.py
View file

@ -1,32 +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 typing import AbstractSet, Dict
from llama_stack.apis.dataset import BaseDataset
class DatasetRegistry:
_REGISTRY: Dict[str, BaseDataset] = {}
@staticmethod
def names() -> AbstractSet[str]:
return DatasetRegistry._REGISTRY.keys()
@staticmethod
def register(name: str, task: BaseDataset) -> None:
if name in DatasetRegistry._REGISTRY:
raise ValueError(f"Dataset {name} already exists.")
DatasetRegistry._REGISTRY[name] = task
@staticmethod
def get_dataset(name: str) -> BaseDataset:
if name not in DatasetRegistry._REGISTRY:
raise ValueError(f"Dataset {name} not found.")
return DatasetRegistry._REGISTRY[name]
@staticmethod
def reset() -> None:
DatasetRegistry._REGISTRY = {}