From 552b747136b2998d0f00490f668d606e15e43039 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 19 Sep 2024 16:31:47 -0700 Subject: [PATCH] dummy model spec --- .../providers/impls/builtin/models/models.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/llama_stack/providers/impls/builtin/models/models.py b/llama_stack/providers/impls/builtin/models/models.py index f87aab83a..08dea21e0 100644 --- a/llama_stack/providers/impls/builtin/models/models.py +++ b/llama_stack/providers/impls/builtin/models/models.py @@ -17,8 +17,13 @@ from llama_models.sku_list import resolve_model from .config import BuiltinImplConfig -DUMMY_MODELS_SPEC = ModelSpec( - llama_model_metadata=resolve_model("Meta-Llama3.1-8B"), +DUMMY_MODELS_SPEC_1 = ModelSpec( + llama_model_metadata=resolve_model("Llama-Guard-3-8B"), + providers_spec={"safety": {"provider_type": "meta-reference"}}, +) + +DUMMY_MODELS_SPEC_2 = ModelSpec( + llama_model_metadata=resolve_model("Meta-Llama3.1-8B-Instruct"), providers_spec={"inference": {"provider_type": "meta-reference"}}, ) @@ -29,16 +34,21 @@ class BuiltinModelsImpl(Models): config: BuiltinImplConfig, ) -> None: self.config = config - self.models_list = [DUMMY_MODELS_SPEC] + self.models = { + x.llama_model_metadata.core_model_id.value: x + for x in [DUMMY_MODELS_SPEC_1, DUMMY_MODELS_SPEC_2] + } async def initialize(self) -> None: pass async def list_models(self) -> ModelsListResponse: - return ModelsListResponse(models_list=self.models_list) + return ModelsListResponse(models_list=list(self.models.values())) async def get_model(self, core_model_id: str) -> ModelsGetResponse: - return ModelsGetResponse(core_model_spec=DUMMY_MODELS_SPEC) + if core_model_id in self.models: + return ModelsGetResponse(core_model_spec=self.models[core_model_id]) + raise ValueError(f"Cannot find {core_model_id} in model registry") async def register_model( self, model_id: str, api: str, provider_spec: Dict[str, str]