diff --git a/pyproject.toml b/pyproject.toml
index e4ff436b..938b43c2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,19 +72,19 @@ all = [
     "openllm[falcon]",
     "openllm[mpt]",
     "openllm[starcoder]",
-    "openllm[fine-tune]",
+    "openllm[ggml]",
     "openllm[openai]",
+    "openllm[agents]",
     "openllm[playground]",
     "openllm[flan-t5]",
-    "openllm[agents]",
-    "openllm[ggml]",
+    "openllm[fine-tune]",
 ]
 chatglm = ["cpm-kernels", "sentencepiece"]
 falcon = ["einops", "xformers", "safetensors"]
 fine-tune = ["peft", "bitsandbytes", "datasets", "accelerate", "deepspeed", "trl"]
 flan-t5 = ["flax", "jax", "jaxlib", "tensorflow", "keras"]
 ggml = ["ctransformers"]
-mpt = ["triton"]
+mpt = ["triton", "einops"]
 openai = ["openai", "tiktoken"]
 playground = ["jupyter", "notebook", "ipython", "jupytext", "nbformat"]
 starcoder = ["bitsandbytes"]
diff --git a/src/openllm/cli.py b/src/openllm/cli.py
index ddebc1a4..6ec279d8 100644
--- a/src/openllm/cli.py
+++ b/src/openllm/cli.py
@@ -2056,9 +2056,8 @@ def prune_command(yes: bool, model_store: ModelStore = Provide[BentoMLContainer.
     """Remove all saved models locally."""
     available = [
         m
-        for t in map(inflection.dasherize, openllm.CONFIG_MAPPING.keys())
         for m in bentoml.models.list()
-        if t in m.tag.name
+        if 'framework' in m.info.labels and m.info.labels['framework'] == 'openllm'
     ]
 
     for model in available:
diff --git a/src/openllm/utils/codegen.py b/src/openllm/utils/codegen.py
index 27fa6593..56d0ee75 100644
--- a/src/openllm/utils/codegen.py
+++ b/src/openllm/utils/codegen.py
@@ -230,6 +230,8 @@ def make_attr_tuple_class(cls_name: str, attr_names: t.Sequence[str]):
         __slots__ = ()
         x = property(itemgetter(0))
     """
+    from . import SHOW_CODEGEN
+
     attr_class_name = f"{cls_name}Attributes"
     attr_class_template = [
         f"class {attr_class_name}(tuple):",
@@ -241,7 +243,11 @@ def make_attr_tuple_class(cls_name: str, attr_names: t.Sequence[str]):
     else:
         attr_class_template.append("    pass")
     globs: DictStrAny = {"_attrs_itemgetter": itemgetter, "_attrs_property": property}
+    if SHOW_CODEGEN:
+        logger.info("Generated class for %s:\n\n%s", attr_class_name, "\n".join(attr_class_template))
+
     _compile_and_eval("\n".join(attr_class_template), globs)
+
     return globs[attr_class_name]
 
 
diff --git a/tools/update-optional-dependencies.py b/tools/update-optional-dependencies.py
index 0ef4ac3f..82fff857 100755
--- a/tools/update-optional-dependencies.py
+++ b/tools/update-optional-dependencies.py
@@ -95,7 +95,7 @@ _NIGHTLY_MAPPING: dict[str, Dependencies] = {
 
 FINE_TUNE_DEPS = ["peft", "bitsandbytes", "datasets", "accelerate", "deepspeed", "trl"]
 FLAN_T5_DEPS = ["flax", "jax", "jaxlib", "tensorflow", "keras"]
-MPT_DEPS = ["triton"]
+MPT_DEPS = ["triton", "einops"]
 OPENAI_DEPS = ["openai", "tiktoken"]
 AGENTS_DEPS = ["transformers[agents]>=4.30", "diffusers", "soundfile"]
 FALCON_DEPS = ["einops", "xformers", "safetensors"]