take hugging face repo

This commit is contained in:
Dinesh Yeduguru 2024-11-18 14:57:56 -08:00
parent acf9af841b
commit 8595b2af85
3 changed files with 25 additions and 16 deletions

View file

@ -31,3 +31,8 @@ def supported_inference_models() -> List[str]:
or is_supported_safety_model(m)
)
]
ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR = {
m.huggingface_repo: m.descriptor() for m in all_registered_models()
}

View file

@ -7,11 +7,14 @@
from collections import namedtuple
from typing import List, Optional
from llama_models.datatypes import CoreModelId
from llama_models.sku_list import all_registered_models
from llama_stack.providers.datatypes import Model, ModelsProtocolPrivate
from llama_stack.providers.utils.inference import (
ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR,
)
ModelAlias = namedtuple("ModelAlias", ["provider_model_id", "aliases", "llama_model"])
@ -77,17 +80,18 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
f"Provider model id '{model.provider_resource_id}' is already registered to a different llama model: '{existing_llama_model}'"
)
else:
# Validate that the llama model is a valid CoreModelId
try:
CoreModelId(model.metadata["llama_model"])
except ValueError as err:
if (
model.metadata["llama_model"]
not in ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR
):
raise ValueError(
f"Invalid llama_model '{model.metadata['llama_model']}' specified in metadata. "
f"Must be one of: {', '.join(m.value for m in CoreModelId)}"
) from err
# Register the mapping from provider model id to llama model for future lookups
f"Must be one of: {', '.join(ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR.keys())}"
)
self.provider_id_to_llama_model_map[model.provider_resource_id] = (
model.metadata["llama_model"]
ALL_HUGGINGFACE_REPOS_TO_MODEL_DESCRIPTOR[
model.metadata["llama_model"]
]
)
return model