update: chat app + llama guard guardrail

application_pages/chat_app.py

@@ -29,6 +29,8 @@ def initialize_session_state():
         st.session_state.test_guardrails = False
     if "llm_model" not in st.session_state:
         st.session_state.llm_model = None
+    if "llama_guard_checkpoint_name" not in st.session_state:
+        st.session_state.llama_guard_checkpoint_name = ""
 
 
 def initialize_guardrails():
@@ -89,6 +91,23 @@ def initialize_guardrails():
                     guardrail_name,
                 )(should_anonymize=True)
             )
+        elif guardrail_name == "PromptInjectionLlamaGuardrail":
+            llama_guard_checkpoint_name = st.sidebar.text_input(
+                "Checkpoint Name", value=""
+            )
+            st.session_state.llama_guard_checkpoint_name = llama_guard_checkpoint_name
+            st.session_state.guardrails.append(
+                getattr(
+                    importlib.import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(
+                    checkpoint=(
+                        None
+                        if st.session_state.llama_guard_checkpoint_name == ""
+                        else st.session_state.llama_guard_checkpoint_name
+                    )
+                )
+            )
         else:
             st.session_state.guardrails.append(
                 getattr(

guardrails_genie/guardrails/injection/llama_prompt_guardrail.py

@@ -76,7 +76,7 @@ class PromptInjectionLlamaGuardrail(Guardrail):
         if self.checkpoint is None:
             self._model = AutoModelForSequenceClassification.from_pretrained(
                 self.model_name
-            ).to(self.device)
+            )
         else:
             api = wandb.Api()
             artifact = api.artifact(self.checkpoint.removeprefix("wandb://"))