update: eval ui

application_pages/evaluation_app.py

@@ -1,6 +1,5 @@
 import asyncio
 import os
-import time
 from importlib import import_module
 
 import pandas as pd
@@ -12,212 +11,81 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 from guardrails_genie.metrics import AccuracyMetric
-from guardrails_genie.utils import EvaluationCallManager
 
 
 def initialize_session_state():
     load_dotenv()
+    if "weave_project_name" not in st.session_state:
+        st.session_state.weave_project_name = "guardrails-genie"
     if "uploaded_file" not in st.session_state:
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
-        st.session_state.dataset_name = ""
+        st.session_state.dataset_name = None
     if "preview_in_app" not in st.session_state:
         st.session_state.preview_in_app = False
+    if "is_dataset_published" not in st.session_state:
+        st.session_state.is_dataset_published = False
+    if "publish_dataset_button" not in st.session_state:
+        st.session_state.publish_dataset_button = 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_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
-    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
-    if "llama_guardrail_checkpoint" not in st.session_state:
-        st.session_state.llama_guardrail_checkpoint = None
-
-
-def initialize_guardrail():
-    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))
-                )
-        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 == "PromptInjectionLlamaGuardrail":
-            llama_guardrail_checkpoint = st.sidebar.text_input(
-                "Llama Guardrail Checkpoint",
-                value=None,
-            )
-            st.session_state.llama_guardrail_checkpoint = llama_guardrail_checkpoint
-            if st.session_state.llama_guardrail_checkpoint is not None:
-                st.session_state.guardrails.append(
-                    getattr(
-                        import_module("guardrails_genie.guardrails"),
-                        guardrail_name,
-                    )(checkpoint=st.session_state.llama_guardrail_checkpoint)
-                )
-        else:
-            st.session_state.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()
 st.title(":material/monitoring: Evaluation")
 
+weave_project_name = st.sidebar.text_input(
+    "Weave project name", value=st.session_state.weave_project_name
+)
+st.session_state.weave_project_name = weave_project_name
+if st.session_state.weave_project_name != "":
+    weave.init(project_name=st.session_state.weave_project_name)
+
 uploaded_file = st.sidebar.file_uploader(
     "Upload the evaluation dataset as a CSV file", type="csv"
 )
 st.session_state.uploaded_file = uploaded_file
-dataset_name = st.sidebar.text_input("Evaluation dataset name", value="")
-st.session_state.dataset_name = dataset_name
-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):
-        dataframe = pd.read_csv(st.session_state.uploaded_file)
-        data_list = dataframe.to_dict(orient="records")
 
-        dataset = weave.Dataset(name=st.session_state.dataset_name, rows=data_list)
-        st.session_state.dataset_ref = weave.publish(dataset)
-
-        entity = st.session_state.dataset_ref.entity
-        project = st.session_state.dataset_ref.project
-        dataset_name = st.session_state.dataset_name
-        digest = st.session_state.dataset_ref._digest
-        st.markdown(
-            f"Dataset published to [**Weave**](https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest})"
+if st.session_state.uploaded_file is not None:
+    dataset_name = st.sidebar.text_input("Evaluation dataset name", value=None)
+    st.session_state.dataset_name = dataset_name
+    preview_in_app = st.sidebar.toggle("Preview in app", value=False)
+    st.session_state.preview_in_app = preview_in_app
+    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 != ""
         )
+    ):
+        
+        with st.expander("Evaluation Dataset Preview", expanded=True):
+            dataframe = pd.read_csv(st.session_state.uploaded_file)
+            data_list = dataframe.to_dict(orient="records")
+
+            dataset = weave.Dataset(name=st.session_state.dataset_name, rows=data_list)
+            st.session_state.dataset_ref = weave.publish(dataset)
+
+            entity = st.session_state.dataset_ref.entity
+            project = st.session_state.dataset_ref.project
+            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})"
+            )
 
-        if preview_in_app:
-            st.dataframe(dataframe)
-
-        st.session_state.dataset_previewed = True
-
-if st.session_state.dataset_previewed:
-    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
-
-    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
-            if st.session_state.start_evaluation:
-                evaluation = weave.Evaluation(
-                    dataset=st.session_state.dataset_ref,
-                    scorers=[AccuracyMetric()],
-                    streamlit_mode=True,
-                )
-                with st.expander("Evaluation Results", expanded=True):
-                    evaluation_summary, call = asyncio.run(
-                        evaluation.evaluate.call(
-                            evaluation,
-                            st.session_state.guardrail_manager,
-                            __weave={
-                                "display_name": "Evaluation.evaluate:"
-                                + st.session_state.evaluation_name
-                            },
-                        )
-                    )
-                    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",
+            if preview_in_app:
+                st.dataframe(dataframe.head(20))
+                if len(dataframe) > 20:
+                    st.markdown(
+                        f"⚠️ Dataset is too large to preview in app, please explore in the [**Weave UI**]({dataset_url})"
                     )
-                    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
+        st.session_state.is_dataset_published = True
+    
+    if st.session_state.is_dataset_published:
+        st.write("Maza Ayega")

guardrails_genie/train/llama_guard.py

@@ -3,12 +3,13 @@ import shutil
 from glob import glob
 from typing import Optional
 
+# import torch.optim as optim
+import bitsandbytes.optim as optim
 import plotly.graph_objects as go
 import streamlit as st
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.optim as optim
 from datasets import load_dataset
 from pydantic import BaseModel
 from rich.progress import track
@@ -335,8 +336,8 @@ class LlamaGuardFineTuner:
 
     def train(
         self,
-        batch_size: int = 32,
-        lr: float = 5e-6,
+        batch_size: int = 16,
+        starting_lr: float = 1e-7,
         num_classes: int = 2,
         log_interval: int = 1,
         save_interval: int = 50,
@@ -358,7 +359,7 @@ class LlamaGuardFineTuner:
 
         Args:
             batch_size (int, optional): The number of samples per batch during training.
-            lr (float, optional): The learning rate for the optimizer.
+            starting_lr (float, optional): The starting learning rate for the optimizer.
             num_classes (int, optional): The number of output classes for the classifier.
             log_interval (int, optional): The interval (in batches) at which to log the loss.
             save_interval (int, optional): The interval (in batches) at which to save model checkpoints.
@@ -377,7 +378,7 @@ class LlamaGuardFineTuner:
         wandb.config.dataset_args = self.dataset_args.model_dump()
         wandb.config.model_name = self.model_name
         wandb.config.batch_size = batch_size
-        wandb.config.lr = lr
+        wandb.config.starting_lr = starting_lr
         wandb.config.num_classes = num_classes
         wandb.config.log_interval = log_interval
         wandb.config.save_interval = save_interval
@@ -387,7 +388,16 @@ class LlamaGuardFineTuner:
         self.model.num_labels = num_classes
         self.model = self.model.to(self.device)
         self.model.train()
-        optimizer = optim.AdamW(self.model.parameters(), lr=lr)
+        # optimizer = optim.AdamW(self.model.parameters(), lr=starting_lr)
+        optimizer = optim.Lion(
+            self.model.parameters(), lr=starting_lr, weight_decay=0.01
+        )
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=starting_lr,
+            steps_per_epoch=len(self.train_dataset) // batch_size + 1,
+            epochs=1,
+        )
         data_loader = DataLoader(
             self.train_dataset,
             batch_size=batch_size,
@@ -405,9 +415,14 @@ class LlamaGuardFineTuner:
             loss = outputs.loss
             optimizer.zero_grad()
             loss.backward()
+
+            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), gradient_clipping)
+
             optimizer.step()
+            scheduler.step()
             if (i + 1) % log_interval == 0:
                 wandb.log({"loss": loss.item()}, step=i + 1)
+                wandb.log({"learning_rate": scheduler.get_last_lr()[0]}, step=i + 1)
             if progress_bar:
                 progress_percentage = (i + 1) * 100 // len(data_loader)
                 progress_bar.progress(