update: evaluation app

application_pages/evaluation_app.py

@@ -1,16 +1,12 @@
-import asyncio
-import os
 from importlib import import_module
 
 import pandas as pd
-import rich
 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
 
 
 def initialize_session_state():
@@ -29,6 +25,99 @@ def initialize_session_state():
         st.session_state.publish_dataset_button = False
     if "dataset_ref" not in st.session_state:
         st.session_state.dataset_ref = None
+    if "guardrails" not in st.session_state:
+        st.session_state.guardrails = []
+    if "guardrail_names" not in st.session_state:
+        st.session_state.guardrail_names = []
+    if "start_evaluations_button" not in st.session_state:
+        st.session_state.start_evaluations_button = False
+
+
+def initialize_guardrails():
+    st.session_state.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:
+                st.session_state.guardrails.append(
+                    getattr(
+                        import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
+                )
+        elif guardrail_name == "PromptInjectionClassifierGuardrail":
+            classifier_model_name = st.sidebar.selectbox(
+                "Classifier Guardrail Model",
+                [
+                    "",
+                    "ProtectAI/deberta-v3-base-prompt-injection-v2",
+                    "wandb://geekyrakshit/guardrails-genie/model-6rwqup9b:v3",
+                ],
+            )
+            if classifier_model_name != "":
+                st.session_state.guardrails.append(
+                    getattr(
+                        import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(model_name=classifier_model_name)
+                )
+        elif guardrail_name == "PresidioEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "RegexEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "TransformersEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "RestrictedTermsJudge":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.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(
+                    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(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )()
+            )
+    st.session_state.guardrails_manager = GuardrailManager(
+        guardrails=st.session_state.guardrails
+    )
 
 
 initialize_session_state()
@@ -54,14 +143,11 @@ if st.session_state.uploaded_file is not None:
     publish_dataset_button = st.sidebar.button("Publish dataset")
     st.session_state.publish_dataset_button = publish_dataset_button
 
-    if (
-        st.session_state.publish_dataset_button
-        and (
-            st.session_state.dataset_name is not None
-            and st.session_state.dataset_name != ""
-        )
+    if st.session_state.publish_dataset_button and (
+        st.session_state.dataset_name is not None
+        and st.session_state.dataset_name != ""
     ):
-        
+
         with st.expander("Evaluation Dataset Preview", expanded=True):
             dataframe = pd.read_csv(st.session_state.uploaded_file)
             data_list = dataframe.to_dict(orient="records")
@@ -74,9 +160,7 @@ if st.session_state.uploaded_file is not None:
             dataset_name = st.session_state.dataset_name
             digest = st.session_state.dataset_ref._digest
             dataset_url = f"https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest}"
-            st.markdown(
-                f"Dataset published to [**Weave**]({dataset_url})"
-            )
+            st.markdown(f"Dataset published to [**Weave**]({dataset_url})")
 
             if preview_in_app:
                 st.dataframe(dataframe.head(20))
@@ -86,6 +170,23 @@ if st.session_state.uploaded_file is not None:
                     )
 
         st.session_state.is_dataset_published = True
-    
+
     if st.session_state.is_dataset_published:
-        st.write("Maza Ayega")
+        guardrail_names = st.sidebar.multiselect(
+            "Select Guardrails",
+            options=[
+                cls_name
+                for cls_name, cls_obj in vars(
+                    import_module("guardrails_genie.guardrails")
+                ).items()
+                if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
+            ],
+        )
+        st.session_state.guardrail_names = guardrail_names
+
+        initialize_guardrails()
+
+        start_evaluations_button = st.sidebar.button("Start Evaluations")
+        st.session_state.start_evaluations_button = start_evaluations_button
+        if st.session_state.start_evaluations_button:
+            st.write(len(st.session_state.guardrails))