From 2980a1892060ea7cbeb7c2a91bff68e5197560c0 Mon Sep 17 00:00:00 2001 From: Matthew Farrellee Date: Tue, 19 Nov 2024 12:49:14 -0500 Subject: [PATCH] map llama model -> provider model id in ModelRegistryHelper --- llama_stack/providers/utils/inference/model_registry.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/utils/inference/model_registry.py b/llama_stack/providers/utils/inference/model_registry.py index 07225fac0..8dbfab14a 100644 --- a/llama_stack/providers/utils/inference/model_registry.py +++ b/llama_stack/providers/utils/inference/model_registry.py @@ -29,7 +29,6 @@ def build_model_alias(provider_model_id: str, model_descriptor: str) -> ModelAli return ModelAlias( provider_model_id=provider_model_id, aliases=[ - model_descriptor, get_huggingface_repo(model_descriptor), ], llama_model=model_descriptor, @@ -57,6 +56,10 @@ class ModelRegistryHelper(ModelsProtocolPrivate): self.alias_to_provider_id_map[alias_obj.provider_model_id] = ( alias_obj.provider_model_id ) + # ensure we can go from llama model to provider model id + self.alias_to_provider_id_map[alias_obj.llama_model] = ( + alias_obj.provider_model_id + ) self.provider_id_to_llama_model_map[alias_obj.provider_model_id] = ( alias_obj.llama_model )