update: LlamaGuardFineTuner

.gitignore

@@ -168,4 +168,5 @@ temp.txt
 binary-classifier/
 wandb/
 artifacts/
-evaluation_results/
+evaluation_results/
+checkpoints/

application_pages/llama_guard_fine_tuning.py

@@ -1,10 +1,16 @@
+import os
+
 import streamlit as st
 
 from guardrails_genie.train.llama_guard import DatasetArgs, LlamaGuardFineTuner
 
 
 def initialize_session_state():
-    st.session_state.llama_guard_fine_tuner = LlamaGuardFineTuner(streamlit_mode=True)
+    st.session_state.llama_guard_fine_tuner = LlamaGuardFineTuner(
+        wandb_project=os.getenv("WANDB_PROJECT_NAME"),
+        wandb_entity=os.getenv("WANDB_ENTITY_NAME"),
+        streamlit_mode=True,
+    )
     if "dataset_address" not in st.session_state:
         st.session_state.dataset_address = ""
     if "train_dataset_range" not in st.session_state:
@@ -25,6 +31,14 @@ def initialize_session_state():
         st.session_state.evaluation_batch_size = None
     if "evaluation_temperature" not in st.session_state:
         st.session_state.evaluation_temperature = None
+    if "checkpoint" not in st.session_state:
+        st.session_state.checkpoint = None
+    if "eval_batch_size" not in st.session_state:
+        st.session_state.eval_batch_size = 32
+    if "eval_positive_label" not in st.session_state:
+        st.session_state.eval_positive_label = 2
+    if "eval_temperature" not in st.session_state:
+        st.session_state.eval_temperature = 1.0
 
 
 initialize_session_state()
@@ -43,18 +57,34 @@ if st.session_state.dataset_address != "":
     st.session_state.train_dataset_range = train_dataset_range
     st.session_state.test_dataset_range = test_dataset_range
 
-    model_name = st.sidebar.selectbox(
-        "Model Name",
-        ["meta-llama/Prompt-Guard-86M"],
+    model_name = st.sidebar.text_input(
+        label="Model Name", value="meta-llama/Prompt-Guard-86M"
     )
     st.session_state.model_name = model_name
 
+    checkpoint = st.sidebar.text_input(label="Fine-tuned Model Checkpoint", value="")
+    st.session_state.checkpoint = checkpoint
+
     preview_dataset = st.sidebar.toggle("Preview Dataset")
     st.session_state.preview_dataset = preview_dataset
 
     evaluate_model = st.sidebar.toggle("Evaluate Model")
     st.session_state.evaluate_model = evaluate_model
 
+    if st.session_state.evaluate_model:
+        eval_batch_size = st.sidebar.slider(
+            label="Eval Batch Size", min_value=16, max_value=1024, value=32
+        )
+        st.session_state.eval_batch_size = eval_batch_size
+
+        eval_positive_label = st.sidebar.number_input("EVal Positive Label", value=2)
+        st.session_state.eval_positive_label = eval_positive_label
+
+        eval_temperature = st.sidebar.slider(
+            label="Eval Temperature", min_value=0.0, max_value=5.0, value=1.0
+        )
+        st.session_state.eval_temperature = eval_temperature
+
     load_fine_tuner_button = st.sidebar.button("Load Fine-Tuner")
     st.session_state.load_fine_tuner_button = load_fine_tuner_button
 
@@ -68,13 +98,19 @@ if st.session_state.dataset_address != "":
                 )
             )
             st.session_state.llama_guard_fine_tuner.load_model(
-                model_name=st.session_state.model_name
+                model_name=st.session_state.model_name,
+                checkpoint=(
+                    None
+                    if st.session_state.checkpoint == ""
+                    else st.session_state.checkpoint
+                ),
             )
             if st.session_state.preview_dataset:
                 st.session_state.llama_guard_fine_tuner.show_dataset_sample()
             if st.session_state.evaluate_model:
                 st.session_state.llama_guard_fine_tuner.evaluate_model(
-                    batch_size=32,
-                    temperature=3.0,
+                    batch_size=st.session_state.eval_batch_size,
+                    positive_label=st.session_state.eval_positive_label,
+                    temperature=st.session_state.eval_temperature,
                 )
             st.session_state.is_fine_tuner_loaded = True

guardrails_genie/train/__init__.py

@@ -1,4 +1,4 @@
+from .llama_guard import DatasetArgs, LlamaGuardFineTuner
 from .train_classifier import train_binary_classifier
-from .llama_guard import LlamaGuardFineTuner, DatasetArgs
 
-__all__ = ["train_binary_classifier", "LlamaGuardFineTuner", "DatasetArgs"]
+__all__ = ["train_binary_classifier", "LlamaGuardFineTuner", "DatasetArgs"]

guardrails_genie/train/llama_guard.py

@@ -1,5 +1,7 @@
 import os
 import shutil
+from glob import glob
+from typing import Optional
 
 import plotly.graph_objects as go
 import streamlit as st
@@ -7,15 +9,16 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
-import wandb
 from datasets import load_dataset
 from pydantic import BaseModel
 from rich.progress import track
-from safetensors.torch import save_model
+from safetensors.torch import load_model, save_model
 from sklearn.metrics import roc_auc_score, roc_curve
 from torch.utils.data import DataLoader
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
+import wandb
+
 
 class DatasetArgs(BaseModel):
     dataset_address: str
@@ -30,7 +33,7 @@ class LlamaGuardFineTuner:
     classification tasks, specifically for detecting prompt injection attacks. It
     integrates with Weights & Biases for experiment tracking and optionally
     displays progress in a Streamlit app.
-    
+
     !!! example "Sample Usage"
         ```python
         from guardrails_genie.train.llama_guard import LlamaGuardFineTuner, DatasetArgs
@@ -98,7 +101,11 @@ class LlamaGuardFineTuner:
             else dataset["test"].select(range(dataset_args.test_dataset_range))
         )
 
-    def load_model(self, model_name: str = "meta-llama/Prompt-Guard-86M"):
+    def load_model(
+        self,
+        model_name: str = "meta-llama/Prompt-Guard-86M",
+        checkpoint: Optional[str] = None,
+    ):
         """
         Loads the specified pre-trained model and tokenizer for sequence classification tasks.
 
@@ -118,9 +125,20 @@ class LlamaGuardFineTuner:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model_name = model_name
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(
-            self.device
-        )
+        if checkpoint is None:
+            self.model = AutoModelForSequenceClassification.from_pretrained(
+                model_name
+            ).to(self.device)
+        else:
+            api = wandb.Api()
+            artifact = api.artifact(checkpoint.removeprefix("wandb://"))
+            artifact_dir = artifact.download()
+            model_file_path = glob(os.path.join(artifact_dir, "model-*.safetensors"))[0]
+            self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+            self.model.classifier = nn.Linear(self.model.classifier.in_features, 2)
+            self.model.num_labels = 2
+            load_model(self.model, model_file_path)
+            self.model = self.model.to(self.device)
 
     def show_dataset_sample(self):
         """

guardrails_genie/train/train_classifier.py

@@ -1,7 +1,6 @@
 import evaluate
 import numpy as np
 import streamlit as st
-import wandb
 from datasets import load_dataset
 from transformers import (
     AutoModelForSequenceClassification,
@@ -11,6 +10,7 @@ from transformers import (
     TrainingArguments,
 )
 
+import wandb
 from guardrails_genie.utils import StreamlitProgressbarCallback