fix(cli): correct prune based on metadata

Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
aarnphm-ec2-dev
2023-07-11 00:34:22 +00:00
parent c2bb29b4f3
commit cea082e7bd
2 changed files with 4 additions and 6 deletions

View File

@@ -72,12 +72,12 @@ all = [
"openllm[falcon]",
"openllm[mpt]",
"openllm[starcoder]",
"openllm[ggml]",
"openllm[openai]",
"openllm[agents]",
"openllm[playground]",
"openllm[flan-t5]",
"openllm[fine-tune]",
"openllm[agents]",
"openllm[playground]",
"openllm[ggml]",
]
chatglm = ["cpm-kernels", "sentencepiece"]
falcon = ["einops", "xformers", "safetensors"]

View File

@@ -2055,9 +2055,7 @@ def models_command(
def prune_command(yes: bool, model_store: ModelStore = Provide[BentoMLContainer.model_store]):
"""Remove all saved models locally."""
available = [
m
for m in bentoml.models.list()
if 'framework' in m.info.labels and m.info.labels['framework'] == 'openllm'
m for m in bentoml.models.list() if "framework" in m.info.labels and m.info.labels["framework"] == "openllm"
]
for model in available: