diff --git a/llama_stack/cli/model/list.py b/llama_stack/cli/model/list.py index 622a6b4e7..b9499f06d 100644 --- a/llama_stack/cli/model/list.py +++ b/llama_stack/cli/model/list.py @@ -75,6 +75,13 @@ class ModelList(Subcommand): action="store_true", help="List the downloaded models", ) + self.parser.add_argument( + "-s", + "--search", + type=str, + required=False, + help="Search for the input string as a substring in the model descriptor(ID)", + ) def _run_model_list_cmd(self, args: argparse.Namespace) -> None: from .safety_models import prompt_guard_model_sku @@ -94,15 +101,19 @@ class ModelList(Subcommand): continue descriptor = model.descriptor() - rows.append( - [ - descriptor, - model.huggingface_repo, - f"{model.max_seq_length // 1024}K", - ] + if not args.search or args.search.lower() in descriptor.lower(): + rows.append( + [ + descriptor, + model.huggingface_repo, + f"{model.max_seq_length // 1024}K", + ] + ) + if len(rows) == 0: + print(f"Did not find any model matching `{args.search}`.") + else: + print_table( + rows, + headers, + separate_rows=True, ) - print_table( - rows, - headers, - separate_rows=True, - )