Make Safety test work, other cleanup

This commit is contained in:
Ashwin Bharambe 2024-10-09 21:09:50 -07:00
parent ba1f294cc6
commit fcd22b6baa
16 changed files with 229 additions and 123 deletions

View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Dict, List, Optional
from typing import Dict, List
from llama_models.sku_list import resolve_model
@ -39,9 +39,3 @@ class ModelRegistryHelper(ModelsProtocolPrivate):
for llama_model, provider_model in self.stack_to_provider_models_map.items():
models.append(ModelDef(identifier=llama_model, llama_model=llama_model))
return models
async def get_model(self, identifier: str) -> Optional[ModelDef]:
if identifier not in self.stack_to_provider_models_map:
return None
return ModelDef(identifier=identifier, llama_model=identifier)