Merge pull request #4 from soumik12345/feat/eval-table

.gitignore

@@ -165,4 +165,7 @@ cursor_prompts/
 uv.lock
 test.py
 temp.txt
-**.csv
+**.csv
+binary-classifier/
+wandb/
+artifacts/

README.md

@@ -18,7 +18,12 @@ source .venv/bin/activate
 ## Run the App
 
 ```bash
-OPENAI_API_KEY="YOUR_OPENAI_API_KEY" streamlit run app.py
+export OPENAI_API_KEY="YOUR_OPENAI_API_KEY"
+export WEAVE_PROJECT="YOUR_WEAVE_PROJECT"
+export WANDB_PROJECT_NAME="YOUR_WANDB_PROJECT_NAME"
+export WANDB_ENTITY_NAME="YOUR_WANDB_ENTITY_NAME"
+export WANDB_LOG_MODEL="checkpoint"
+streamlit run app.py
 ```
 
 ## Use the Library

app.py

@@ -13,6 +13,13 @@ evaluation_page = st.Page(
     title="Evaluation",
     icon=":material/monitoring:",
 )
-page_navigation = st.navigation([intro_page, chat_page, evaluation_page])
+train_classifier_page = st.Page(
+    "application_pages/train_classifier.py",
+    title="Train Classifier",
+    icon=":material/fitness_center:",
+)
+page_navigation = st.navigation(
+    [intro_page, chat_page, evaluation_page, train_classifier_page]
+)
 st.set_page_config(page_title="Guardrails Genie", page_icon=":material/guardian:")
 page_navigation.run()

application_pages/chat_app.py

@@ -1,4 +1,5 @@
 import importlib
+import os
 
 import streamlit as st
 import weave
@@ -7,27 +8,27 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 
-st.title(":material/robot: Guardrails Genie Playground")
-
-load_dotenv()
-weave.init(project_name="guardrails-genie")
 
-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 "guardrails_manager" not in st.session_state:
-    st.session_state.guardrails_manager = None
-if "initialize_guardrails" not in st.session_state:
-    st.session_state.initialize_guardrails = False
-if "system_prompt" not in st.session_state:
-    st.session_state.system_prompt = ""
-if "user_prompt" not in st.session_state:
-    st.session_state.user_prompt = ""
-if "test_guardrails" not in st.session_state:
-    st.session_state.test_guardrails = False
-if "llm_model" not in st.session_state:
-    st.session_state.llm_model = None
+def initialize_session_state():
+    load_dotenv()
+    weave.init(project_name=os.getenv("WEAVE_PROJECT"))
+
+    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 "guardrails_manager" not in st.session_state:
+        st.session_state.guardrails_manager = None
+    if "initialize_guardrails" not in st.session_state:
+        st.session_state.initialize_guardrails = False
+    if "system_prompt" not in st.session_state:
+        st.session_state.system_prompt = ""
+    if "user_prompt" not in st.session_state:
+        st.session_state.user_prompt = ""
+    if "test_guardrails" not in st.session_state:
+        st.session_state.test_guardrails = False
+    if "llm_model" not in st.session_state:
+        st.session_state.llm_model = None
 
 
 def initialize_guardrails():
@@ -44,18 +45,30 @@ def initialize_guardrails():
                         guardrail_name,
                     )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
                 )
-        else:
-            st.session_state.guardrails.append(
-                getattr(
-                    importlib.import_module("guardrails_genie.guardrails"),
-                    guardrail_name,
-                )()
+        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(
+                        importlib.import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(model_name=classifier_model_name)
+                )
     st.session_state.guardrails_manager = GuardrailManager(
         guardrails=st.session_state.guardrails
     )
 
 
+initialize_session_state()
+st.title(":material/robot: Guardrails Genie Playground")
+
 openai_model = st.sidebar.selectbox(
     "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
 )
@@ -97,7 +110,7 @@ if st.session_state.initialize_guardrails:
 
         if guardrails_response["safe"]:
             st.markdown(
-                f"\n\n---\nPrompt is safe! Explore prompt trace on [Weave]({call.ui_url})\n\n---\n"
+                f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
             )
 
             with st.sidebar.status("Generating response from LLM..."):

application_pages/evaluation_app.py

@@ -1,7 +1,10 @@
 import asyncio
+import os
+import time
 from importlib import import_module
 
 import pandas as pd
+import rich
 import streamlit as st
 import weave
 from dotenv import load_dotenv
@@ -9,12 +12,11 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 from guardrails_genie.metrics import AccuracyMetric
-
-load_dotenv()
-weave.init(project_name="guardrails-genie")
+from guardrails_genie.utils import EvaluationCallManager
 
 
 def initialize_session_state():
+    load_dotenv()
     if "uploaded_file" not in st.session_state:
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
@@ -35,6 +37,18 @@ def initialize_session_state():
         st.session_state.evaluation_summary = None
     if "guardrail_manager" not in st.session_state:
         st.session_state.guardrail_manager = None
+    if "evaluation_name" not in st.session_state:
+        st.session_state.evaluation_name = ""
+    if "show_result_table" not in st.session_state:
+        st.session_state.show_result_table = False
+    if "weave_client" not in st.session_state:
+        st.session_state.weave_client = weave.init(
+            project_name=os.getenv("WEAVE_PROJECT")
+        )
+    if "evaluation_call_manager" not in st.session_state:
+        st.session_state.evaluation_call_manager = None
+    if "call_id" not in st.session_state:
+        st.session_state.call_id = None
 
 
 def initialize_guardrail():
@@ -51,10 +65,22 @@ def initialize_guardrail():
                         guardrail_name,
                     )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
                 )
-        else:
-            guardrails.append(
-                getattr(import_module("guardrails_genie.guardrails"), guardrail_name)()
+        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)
+                )
     st.session_state.guardrails = guardrails
     st.session_state.guardrail_manager = GuardrailManager(guardrails=guardrails)
 
@@ -107,6 +133,8 @@ if st.session_state.dataset_previewed:
 
     if st.session_state.guardrail_names != []:
         initialize_guardrail()
+        evaluation_name = st.sidebar.text_input("Evaluation name", value="")
+        st.session_state.evaluation_name = evaluation_name
         if st.session_state.guardrail_manager is not None:
             if st.sidebar.button("Start Evaluation"):
                 st.session_state.start_evaluation = True
@@ -119,10 +147,55 @@ if st.session_state.dataset_previewed:
                 with st.expander("Evaluation Results", expanded=True):
                     evaluation_summary, call = asyncio.run(
                         evaluation.evaluate.call(
-                            evaluation, st.session_state.guardrail_manager
+                            evaluation,
+                            st.session_state.guardrail_manager,
+                            __weave={
+                                "display_name": "Evaluation.evaluate:"
+                                + st.session_state.evaluation_name
+                            },
                         )
                     )
-                    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
+                    x_axis = list(evaluation_summary["AccuracyMetric"].keys())
+                    y_axis = [
+                        evaluation_summary["AccuracyMetric"][x_axis_item]
+                        for x_axis_item in x_axis
+                    ]
+                    st.bar_chart(
+                        pd.DataFrame({"Metric": x_axis, "Score": y_axis}),
+                        x="Metric",
+                        y="Score",
+                    )
+                    st.session_state.evaluation_summary = evaluation_summary
+                    st.session_state.call_id = call.id
+                    st.session_state.start_evaluation = False
+
+                    if not st.session_state.start_evaluation:
+                        time.sleep(5)
+                        st.session_state.evaluation_call_manager = (
+                            EvaluationCallManager(
+                                entity="geekyrakshit",
+                                project="guardrails-genie",
+                                call_id=st.session_state.call_id,
+                            )
+                        )
+                        for guardrail_name in st.session_state.guardrail_names:
+                            st.session_state.evaluation_call_manager.call_list.append(
+                                {
+                                    "guardrail_name": guardrail_name,
+                                    "calls": st.session_state.evaluation_call_manager.collect_guardrail_guard_calls_from_eval(),
+                                }
+                            )
+                            rich.print(
+                                st.session_state.evaluation_call_manager.call_list
+                            )
+                        st.dataframe(
+                            st.session_state.evaluation_call_manager.render_calls_to_streamlit()
+                        )
+                        if st.session_state.evaluation_call_manager.show_warning_in_app:
+                            st.warning(
+                                f"Only {st.session_state.evaluation_call_manager.max_count} calls can be shown in the app."
+                            )
+                        st.markdown(
+                            f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
+                        )
+                    st.session_state.evaluation_call_manager = None

application_pages/train_classifier.py

@@ -0,0 +1,61 @@
+import os
+
+import streamlit as st
+from dotenv import load_dotenv
+
+from guardrails_genie.train_classifier import train_binary_classifier
+
+
+def initialize_session_state():
+    load_dotenv()
+    if "dataset_name" not in st.session_state:
+        st.session_state.dataset_name = None
+    if "base_model_name" not in st.session_state:
+        st.session_state.base_model_name = None
+    if "batch_size" not in st.session_state:
+        st.session_state.batch_size = 16
+    if "should_start_training" not in st.session_state:
+        st.session_state.should_start_training = False
+    if "training_output" not in st.session_state:
+        st.session_state.training_output = None
+
+
+initialize_session_state()
+st.title(":material/fitness_center: Train Classifier")
+
+dataset_name = st.sidebar.text_input("Dataset Name", value="")
+st.session_state.dataset_name = dataset_name
+
+base_model_name = st.sidebar.selectbox(
+    "Base Model",
+    options=[
+        "distilbert/distilbert-base-uncased",
+        "FacebookAI/roberta-base",
+        "microsoft/deberta-v3-base",
+    ],
+)
+st.session_state.base_model_name = base_model_name
+
+batch_size = st.sidebar.slider(
+    "Batch Size", min_value=4, max_value=256, value=16, step=4
+)
+st.session_state.batch_size = batch_size
+
+train_button = st.sidebar.button("Train")
+st.session_state.should_start_training = (
+    train_button and st.session_state.dataset_name and st.session_state.base_model_name
+)
+
+if st.session_state.should_start_training:
+    with st.expander("Training", expanded=True):
+        training_output = train_binary_classifier(
+            project_name=os.getenv("WANDB_PROJECT_NAME"),
+            entity_name=os.getenv("WANDB_ENTITY_NAME"),
+            run_name=f"{st.session_state.base_model_name}-finetuned",
+            dataset_repo=st.session_state.dataset_name,
+            model_name=st.session_state.base_model_name,
+            batch_size=st.session_state.batch_size,
+            streamlit_mode=True,
+        )
+        st.session_state.training_output = training_output
+        st.write(training_output)

guardrails_genie/guardrails/__init__.py

@@ -1,8 +1,11 @@
-from .injection import PromptInjectionProtectAIGuardrail, PromptInjectionSurveyGuardrail
+from .injection import (
+    PromptInjectionClassifierGuardrail,
+    PromptInjectionSurveyGuardrail,
+)
 from .manager import GuardrailManager
 
 __all__ = [
     "PromptInjectionSurveyGuardrail",
-    "PromptInjectionProtectAIGuardrail",
+    "PromptInjectionClassifierGuardrail",
     "GuardrailManager",
 ]

guardrails_genie/guardrails/injection/__init__.py

@@ -1,4 +1,4 @@
-from .protectai_guardrail import PromptInjectionProtectAIGuardrail
+from .classifier_guardrail import PromptInjectionClassifierGuardrail
 from .survey_guardrail import PromptInjectionSurveyGuardrail
 
-__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionProtectAIGuardrail"]
+__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionClassifierGuardrail"]

guardrails_genie/guardrails/injection/{protectai_guardrail.py → classifier_guardrail.py}

@@ -5,16 +5,25 @@ import weave
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 from transformers.pipelines.base import Pipeline
 
+import wandb
+
 from ..base import Guardrail
 
 
-class PromptInjectionProtectAIGuardrail(Guardrail):
+class PromptInjectionClassifierGuardrail(Guardrail):
     model_name: str = "ProtectAI/deberta-v3-base-prompt-injection-v2"
     _classifier: Optional[Pipeline] = None
 
     def model_post_init(self, __context):
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+        if self.model_name.startswith("wandb://"):
+            api = wandb.Api()
+            artifact = api.artifact(self.model_name.removeprefix("wandb://"))
+            artifact_dir = artifact.download()
+            tokenizer = AutoTokenizer.from_pretrained(artifact_dir)
+            model = AutoModelForSequenceClassification.from_pretrained(artifact_dir)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
         self._classifier = pipeline(
             "text-classification",
             model=model,
@@ -28,11 +37,6 @@ class PromptInjectionProtectAIGuardrail(Guardrail):
     def classify(self, prompt: str):
         return self._classifier(prompt)
 
-    @weave.op()
-    def predict(self, prompt: str):
-        response = self.classify(prompt)
-        return {"safe": response[0]["label"] != "INJECTION"}
-
     @weave.op()
     def guard(self, prompt: str):
         response = self.classify(prompt)
@@ -41,3 +45,7 @@ class PromptInjectionProtectAIGuardrail(Guardrail):
             "safe": response[0]["label"] != "INJECTION",
             "summary": f"Prompt is deemed {response[0]['label']} with {confidence_percentage}% confidence.",
         }
+
+    @weave.op()
+    def predict(self, prompt: str):
+        return self.guard(prompt)

guardrails_genie/train_classifier.py

@@ -0,0 +1,110 @@
+import evaluate
+import numpy as np
+import streamlit as st
+from datasets import load_dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from transformers.trainer_callback import TrainerControl, TrainerState
+
+import wandb
+
+
+class StreamlitProgressbarCallback(TrainerCallback):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.progress_bar = st.progress(0, text="Training")
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        super().on_step_begin(args, state, control, **kwargs)
+        self.progress_bar.progress(
+            (state.global_step * 100 // state.max_steps) + 1,
+            text=f"Training {state.global_step} / {state.max_steps}",
+        )
+
+
+def train_binary_classifier(
+    project_name: str,
+    entity_name: str,
+    run_name: str,
+    dataset_repo: str = "geekyrakshit/prompt-injection-dataset",
+    model_name: str = "distilbert/distilbert-base-uncased",
+    learning_rate: float = 2e-5,
+    batch_size: int = 16,
+    num_epochs: int = 2,
+    weight_decay: float = 0.01,
+    streamlit_mode: bool = False,
+):
+    wandb.init(project=project_name, entity=entity_name, name=run_name)
+    if streamlit_mode:
+        st.markdown(
+            f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
+        )
+    dataset = load_dataset(dataset_repo)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def preprocess_function(examples):
+        return tokenizer(examples["prompt"], truncation=True)
+
+    tokenized_datasets = dataset.map(preprocess_function, batched=True)
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    accuracy = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        predictions = np.argmax(predictions, axis=1)
+        return accuracy.compute(predictions=predictions, references=labels)
+
+    id2label = {0: "SAFE", 1: "INJECTION"}
+    label2id = {"SAFE": 0, "INJECTION": 1}
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        num_labels=2,
+        id2label=id2label,
+        label2id=label2id,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=TrainingArguments(
+            output_dir="binary-classifier",
+            learning_rate=learning_rate,
+            per_device_train_batch_size=batch_size,
+            per_device_eval_batch_size=batch_size,
+            num_train_epochs=num_epochs,
+            weight_decay=weight_decay,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            load_best_model_at_end=True,
+            push_to_hub=True,
+            report_to="wandb",
+            logging_strategy="steps",
+            logging_steps=1,
+        ),
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["test"],
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+        callbacks=[StreamlitProgressbarCallback()] if streamlit_mode else [],
+    )
+    try:
+        training_output = trainer.train()
+    except Exception as e:
+        wandb.finish()
+        raise e
+    wandb.finish()
+    return training_output

guardrails_genie/utils.py

@@ -1,7 +1,9 @@
 import os
 
+import pandas as pd
 import pymupdf4llm
 import weave
+import weave.trace
 from firerequests import FireRequests
 
 
@@ -11,3 +13,47 @@ def get_markdown_from_pdf_url(url: str) -> str:
     markdown = pymupdf4llm.to_markdown("temp.pdf", show_progress=False)
     os.remove("temp.pdf")
     return markdown
+
+
+class EvaluationCallManager:
+    def __init__(self, entity: str, project: str, call_id: str, max_count: int = 10):
+        self.base_call = weave.init(f"{entity}/{project}").get_call(call_id=call_id)
+        self.max_count = max_count
+        self.show_warning_in_app = False
+        self.call_list = []
+
+    def collect_guardrail_guard_calls_from_eval(self):
+        guard_calls, count = [], 0
+        for eval_predict_and_score_call in self.base_call.children():
+            if "Evaluation.summarize" in eval_predict_and_score_call._op_name:
+                break
+            guardrail_predict_call = eval_predict_and_score_call.children()[0]
+            guard_call = guardrail_predict_call.children()[0]
+            score_call = eval_predict_and_score_call.children()[1]
+            guard_calls.append(
+                {
+                    "input_prompt": str(guard_call.inputs["prompt"]),
+                    "outputs": dict(guard_call.output),
+                    "score": dict(score_call.output),
+                }
+            )
+            count += 1
+            if count >= self.max_count:
+                self.show_warning_in_app = True
+                break
+        return guard_calls
+
+    def render_calls_to_streamlit(self):
+        dataframe = {
+            "input_prompt": [
+                call["input_prompt"] for call in self.call_list[0]["calls"]
+            ]
+        }
+        for guardrail_call in self.call_list:
+            dataframe[guardrail_call["guardrail_name"] + ".safe"] = [
+                call["outputs"]["safe"] for call in guardrail_call["calls"]
+            ]
+            dataframe[guardrail_call["guardrail_name"] + ".prediction_correctness"] = [
+                call["score"]["correct"] for call in guardrail_call["calls"]
+            ]
+        return pd.DataFrame(dataframe)

pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "ruff>=0.6.9",
     "pip>=24.2",
     "uv>=0.4.20",
-    "git+https://github.com/wandb/weave@feat/eval-progressbar",
+    "weave @ git+https://github.com/wandb/weave@feat/eval-progressbar",
     "streamlit>=1.40.1",
     "python-dotenv>=1.0.1",
     "watchdog>=6.0.0",
@@ -23,4 +23,4 @@ dependencies = [
 ]
 
 [tool.setuptools]
-py-modules = ["guardrails_genie"]
+py-modules = ["guardrails_genie"]