Add special case for prompt guard

llama_stack/cli/model/describe.py

@@ -39,7 +39,14 @@ class ModelDescribe(Subcommand):
         )
 
     def _run_model_describe_cmd(self, args: argparse.Namespace) -> None:
-        model = resolve_model(args.model_id)
+        from .safety_models import prompt_guard_model_sku
+
+        prompt_guard = prompt_guard_model_sku()
+        if args.model_id == prompt_guard.model_id:
+            model = prompt_guard
+        else:
+            model = resolve_model(args.model_id)
+
         if model is None:
             self.parser.error(
                 f"Model {args.model_id} not found; try 'llama model list' for a list of available models."

llama_stack/cli/model/list.py

@@ -34,6 +34,8 @@ class ModelList(Subcommand):
         )
 
     def _run_model_list_cmd(self, args: argparse.Namespace) -> None:
+        from .safety_models import prompt_guard_model_sku
+
         headers = [
             "Model Descriptor",
             "Hugging Face Repo",
@@ -41,7 +43,7 @@ class ModelList(Subcommand):
         ]
 
         rows = []
-        for model in all_registered_models():
+        for model in all_registered_models() + [prompt_guard_model_sku()]:
             if not args.show_all and not model.is_featured:
                 continue
 

llama_stack/cli/model/safety_models.py

@@ -0,0 +1,52 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from llama_models.datatypes import *  # noqa: F403
+from llama_models.sku_list import LlamaDownloadInfo
+
+
+class PromptGuardModel(BaseModel):
+    """Make a 'fake' Model-like object for Prompt Guard. Eventually this will be removed."""
+
+    model_id: str = "Prompt-Guard-86M"
+    description: str = (
+        "Prompt Guard. NOTE: this model will not be provided via `llama` CLI soon."
+    )
+    is_featured: bool = False
+    huggingface_repo: str = "meta-llama/Prompt-Guard-86M"
+    max_seq_length: int = 2048
+    is_instruct_model: bool = False
+    quantization_format: CheckpointQuantizationFormat = (
+        CheckpointQuantizationFormat.bf16
+    )
+    arch_args: Dict[str, Any] = Field(default_factory=dict)
+    recommended_sampling_params: Optional[SamplingParams] = None
+
+    def descriptor(self) -> str:
+        return self.model_id
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+def prompt_guard_model_sku():
+    return PromptGuardModel()
+
+
+def prompt_guard_download_info():
+    return LlamaDownloadInfo(
+        folder="Prompt-Guard",
+        files=[
+            "model.safetensors",
+            "special_tokens_map.json",
+            "tokenizer.json",
+            "tokenizer_config.json",
+        ],
+        pth_size=1,
+    )