浏览代码

Update promptguard model-id

Suraj 9 月之前
父节点
当前提交
d1d08f9b82

+ 1 - 1
recipes/responsible_ai/prompt_guard/Prompt Guard Tutorial.ipynb

@@ -55,7 +55,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "prompt_injection_model_name = 'meta-llama/PromptGuard'\n",
+    "prompt_injection_model_name = 'meta-llama/Prompt-Guard-86M'\n",
     "tokenizer = AutoTokenizer.from_pretrained(prompt_injection_model_name)\n",
     "model = AutoModelForSequenceClassification.from_pretrained(prompt_injection_model_name)"
    ]

+ 2 - 2
recipes/responsible_ai/prompt_guard/inference.py

@@ -11,12 +11,12 @@ Utilities for loading the PromptGuard model and evaluating text for jailbreaks a
 """
 
 
-def load_model_and_tokenizer(model_name='meta-llama/PromptGuard'):
+def load_model_and_tokenizer(model_name='meta-llama/Prompt-Guard-86M'):
     """
     Load the PromptGuard model from Hugging Face or a local model.
     
     Args:
-        model_name (str): The name of the model to load. Default is 'meta-llama/PromptGuard'.
+        model_name (str): The name of the model to load. Default is 'meta-llama/Prompt-Guard-86M'.
         
     Returns:
         transformers.PreTrainedModel: The loaded model.