From cea082e7bd9c0ba2b2bd0185d8c43d2eb2585e4d Mon Sep 17 00:00:00 2001
From: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
Date: Tue, 11 Jul 2023 00:34:22 +0000
Subject: [PATCH] fix(cli): correct prune based on metadata

Signed-off-by: aarnphm-ec2-dev <29749331+aarnphm@users.noreply.github.com>
---
 pyproject.toml     | 6 +++---
 src/openllm/cli.py | 4 +---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 938b43c2..ccee2ca4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"]
diff --git a/src/openllm/cli.py b/src/openllm/cli.py
index 6ec279d8..9e36b15e 100644
--- a/src/openllm/cli.py
+++ b/src/openllm/cli.py
@@ -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: