feat: add substring search option for model list

Signed-off-by: reidliu <reid201711@gmail.com>
This commit is contained in:
reidliu 2025-02-20 14:33:41 +08:00
parent e7d261ef4a
commit ad5ea63140

View file

@ -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,
)