update: evaluation app

application_pages/evaluation_app.py

@@ -6,6 +6,7 @@ import streamlit as st
 import weave
 from dotenv import load_dotenv
 
+from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 from guardrails_genie.metrics import AccuracyMetric
 
@@ -18,37 +19,44 @@ def initialize_session_state():
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
         st.session_state.dataset_name = ""
-    if "visualize_in_app" not in st.session_state:
-        st.session_state.visualize_in_app = False
+    if "preview_in_app" not in st.session_state:
+        st.session_state.preview_in_app = False
     if "dataset_ref" not in st.session_state:
         st.session_state.dataset_ref = None
     if "dataset_previewed" not in st.session_state:
         st.session_state.dataset_previewed = False
-    if "guardrail_name" not in st.session_state:
-        st.session_state.guardrail_name = ""
-    if "guardrail" not in st.session_state:
-        st.session_state.guardrail = None
+    if "guardrail_names" not in st.session_state:
+        st.session_state.guardrail_names = []
+    if "guardrails" not in st.session_state:
+        st.session_state.guardrails = []
     if "start_evaluation" not in st.session_state:
         st.session_state.start_evaluation = False
     if "evaluation_summary" not in st.session_state:
         st.session_state.evaluation_summary = None
+    if "guardrail_manager" not in st.session_state:
+        st.session_state.guardrail_manager = None
 
 
 def initialize_guardrail():
-    if st.session_state.guardrail_name == "PromptInjectionSurveyGuardrail":
-        survey_guardrail_model = st.sidebar.selectbox(
-            "Survey Guardrail LLM", ["", "gpt-4o-mini", "gpt-4o"]
-        )
-        if survey_guardrail_model:
-            st.session_state.guardrail = getattr(
-                import_module("guardrails_genie.guardrails"),
-                st.session_state.guardrail_name,
-            )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
-    else:
-        st.session_state.guardrail = getattr(
-            import_module("guardrails_genie.guardrails"),
-            st.session_state.guardrail_name,
-        )()
+    guardrails = []
+    for guardrail_name in st.session_state.guardrail_names:
+        if guardrail_name == "PromptInjectionSurveyGuardrail":
+            survey_guardrail_model = st.sidebar.selectbox(
+                "Survey Guardrail LLM", ["", "gpt-4o-mini", "gpt-4o"]
+            )
+            if survey_guardrail_model:
+                guardrails.append(
+                    getattr(
+                        import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
+                )
+        else:
+            guardrails.append(
+                getattr(import_module("guardrails_genie.guardrails"), guardrail_name)()
+            )
+    st.session_state.guardrails = guardrails
+    st.session_state.guardrail_manager = GuardrailManager(guardrails=guardrails)
 
 
 initialize_session_state()
@@ -60,8 +68,8 @@ uploaded_file = st.sidebar.file_uploader(
 st.session_state.uploaded_file = uploaded_file
 dataset_name = st.sidebar.text_input("Evaluation dataset name", value="")
 st.session_state.dataset_name = dataset_name
-visualize_in_app = st.sidebar.toggle("Visualize in app", value=False)
-st.session_state.visualize_in_app = visualize_in_app
+preview_in_app = st.sidebar.toggle("Preview in app", value=False)
+st.session_state.preview_in_app = preview_in_app
 
 if st.session_state.uploaded_file is not None and st.session_state.dataset_name != "":
     with st.expander("Evaluation Dataset Preview", expanded=True):
@@ -79,16 +87,15 @@ if st.session_state.uploaded_file is not None and st.session_state.dataset_name
             f"Dataset published to [**Weave**](https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest})"
         )
 
-        if visualize_in_app:
+        if preview_in_app:
             st.dataframe(dataframe)
 
         st.session_state.dataset_previewed = True
 
 if st.session_state.dataset_previewed:
-    guardrail_name = st.sidebar.selectbox(
-        "Select Guardrail",
-        options=[""]
-        + [
+    guardrail_names = st.sidebar.multiselect(
+        "Select Guardrails",
+        options=[
             cls_name
             for cls_name, cls_obj in vars(
                 import_module("guardrails_genie.guardrails")
@@ -96,11 +103,11 @@ if st.session_state.dataset_previewed:
             if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
         ],
     )
-    st.session_state.guardrail_name = guardrail_name
+    st.session_state.guardrail_names = guardrail_names
 
-    if st.session_state.guardrail_name != "":
+    if st.session_state.guardrail_names != []:
         initialize_guardrail()
-        if st.session_state.guardrail is not None:
+        if st.session_state.guardrail_manager is not None:
             if st.sidebar.button("Start Evaluation"):
                 st.session_state.start_evaluation = True
             if st.session_state.start_evaluation:
@@ -110,9 +117,12 @@ if st.session_state.dataset_previewed:
                     streamlit_mode=True,
                 )
                 with st.expander("Evaluation Results", expanded=True):
-                    evaluation_summary = asyncio.run(
-                        evaluation.evaluate(st.session_state.guardrail)
+                    evaluation_summary, call = asyncio.run(
+                        evaluation.evaluate.call(
+                            evaluation, st.session_state.guardrail_manager
+                        )
                     )
+                    st.markdown(f"[Explore evaluation in Weave]({call.ui_url})")
                     st.write(evaluation_summary)
                 st.session_state.evaluation_summary = evaluation_summary
                 st.session_state.start_evaluation = False

guardrails_genie/guardrails/injection/protectai_guardrail.py

@@ -23,7 +23,7 @@ class PromptInjectionProtectAIGuardrail(Guardrail):
             max_length=512,
             device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
         )
-    
+
     @weave.op()
     def classify(self, prompt: str):
         return self._classifier(prompt)

guardrails_genie/guardrails/manager.py

@@ -1,20 +1,28 @@
 import weave
 from rich.progress import track
-from weave.flow.obj import Object as WeaveObject
 
 from .base import Guardrail
 
 
-class GuardrailManager(WeaveObject):
+class GuardrailManager(weave.Model):
     guardrails: list[Guardrail]
 
     @weave.op()
-    def guard(self, prompt: str, **kwargs) -> dict:
+    def guard(self, prompt: str, progress_bar: bool = True, **kwargs) -> dict:
         alerts, safe = [], True
-        for guardrail in track(self.guardrails, description="Running guardrails"):
+        iterable = (
+            track(self.guardrails, description="Running guardrails")
+            if progress_bar
+            else self.guardrails
+        )
+        for guardrail in iterable:
             response = guardrail.guard(prompt, **kwargs)
             alerts.append(
                 {"guardrail_name": guardrail.__class__.__name__, "response": response}
             )
             safe = safe and response["safe"]
         return {"safe": safe, "alerts": alerts}
+
+    @weave.op()
+    def predict(self, prompt: str, **kwargs) -> dict:
+        return self.guard(prompt, progress_bar=False, **kwargs)

guardrails_genie/metrics.py

@@ -17,11 +17,6 @@ class AccuracyMetric(weave.Scorer):
         count_true = list(valid_data).count(True)
         int_data = [int(x) for x in valid_data]
 
-        sample_mean = np.mean(int_data) if int_data else 0
-        sample_variance = np.var(int_data) if int_data else 0
-        sample_error = np.sqrt(sample_variance / len(int_data)) if int_data else 0
-
-        # Calculate precision, recall, and F1 score
         true_positives = count_true
         false_positives = len(valid_data) - count_true
         false_negatives = len(score_rows) - len(valid_data)
@@ -43,7 +38,7 @@ class AccuracyMetric(weave.Scorer):
         )
 
         return {
-            "accuracy": float(sample_mean),
+            "accuracy": float(np.mean(int_data) if int_data else 0),
             "precision": precision,
             "recall": recall,
             "f1_score": f1_score,