add: llama guard fine-tuner[WIP]

app.py

@@ -18,8 +18,19 @@ train_classifier_page = st.Page(
     title="Train Classifier",
     icon=":material/fitness_center:",
 )
+llama_guard_fine_tuning_page = st.Page(
+    "application_pages/llama_guard_fine_tuning.py",
+    title="Fine-Tune LLama Guard",
+    icon=":material/star:",
+)
 page_navigation = st.navigation(
-    [intro_page, chat_page, evaluation_page, train_classifier_page]
+    [
+        intro_page,
+        chat_page,
+        evaluation_page,
+        train_classifier_page,
+        llama_guard_fine_tuning_page,
+    ]
 )
 st.set_page_config(page_title="Guardrails Genie", page_icon=":material/guardian:")
 page_navigation.run()

application_pages/llama_guard_fine_tuning.py

@@ -0,0 +1,43 @@
+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)
+    if "dataset_address" not in st.session_state:
+        st.session_state.dataset_address = ""
+    if "train_dataset_range" not in st.session_state:
+        st.session_state.train_dataset_range = 0
+    if "test_dataset_range" not in st.session_state:
+        st.session_state.test_dataset_range = 0
+    if "load_dataset_button" not in st.session_state:
+        st.session_state.load_dataset_button = False
+
+
+initialize_session_state()
+st.title(":material/star: Fine-Tune LLama Guard")
+
+dataset_address = st.sidebar.text_input("Dataset Address", value="")
+st.session_state.dataset_address = dataset_address
+
+if st.session_state.dataset_address != "":
+    train_dataset_range = st.sidebar.number_input(
+        "Train Dataset Range", value=0, min_value=0, max_value=252956
+    )
+    test_dataset_range = st.sidebar.number_input(
+        "Test Dataset Range", value=0, min_value=0, max_value=63240
+    )
+    st.session_state.train_dataset_range = train_dataset_range
+    st.session_state.test_dataset_range = test_dataset_range
+    load_dataset_button = st.sidebar.button("Load Dataset")
+    st.session_state.load_dataset_button = load_dataset_button
+    if load_dataset_button:
+        with st.status("Dataset Arguments"):
+            dataset_args = DatasetArgs(
+                dataset_address=st.session_state.dataset_address,
+                train_dataset_range=st.session_state.train_dataset_range,
+                test_dataset_range=st.session_state.test_dataset_range,
+            )
+            st.session_state.llama_guard_fine_tuner.load_dataset(dataset_args)
+            st.session_state.llama_guard_fine_tuner.show_dataset_sample()

guardrails_genie/train/llama_guard.py

@@ -0,0 +1,128 @@
+import matplotlib.pyplot as plt
+import streamlit as st
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from pydantic import BaseModel
+from rich.progress import track
+from sklearn.metrics import roc_auc_score, roc_curve
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+
+class DatasetArgs(BaseModel):
+    dataset_address: str
+    train_dataset_range: int
+    test_dataset_range: int
+
+
+class LlamaGuardFineTuner:
+    def __init__(self, streamlit_mode: bool = False):
+        self.streamlit_mode = streamlit_mode
+
+    def load_dataset(self, dataset_args: DatasetArgs):
+        dataset = load_dataset(dataset_args.dataset_address)
+        self.train_dataset = (
+            dataset["train"]
+            if dataset_args.train_dataset_range > 0
+            else dataset["train"].select(range(dataset_args.train_dataset_range))
+        )
+        self.test_dataset = (
+            dataset["test"]
+            if dataset_args.test_dataset_range > 0
+            else dataset["test"].select(range(dataset_args.test_dataset_range))
+        )
+
+    def load_model(self, model_name: str = "meta-llama/Prompt-Guard-86M"):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(
+            self.device
+        )
+
+    def show_dataset_sample(self):
+        if self.streamlit_mode:
+            st.markdown("### Train Dataset Sample")
+            st.dataframe(self.train_dataset.to_pandas().head())
+            st.markdown("### Test Dataset Sample")
+            st.dataframe(self.test_dataset.to_pandas().head())
+
+    def evaluate_batch(
+        self,
+        texts,
+        batch_size: int = 32,
+        positive_label: int = 2,
+        temperature: float = 1.0,
+        truncation: bool = True,
+        max_length: int = 512,
+    ) -> list[float]:
+        self.model.eval()
+        encoded_texts = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        dataset = torch.utils.data.TensorDataset(
+            encoded_texts["input_ids"], encoded_texts["attention_mask"]
+        )
+        data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
+
+        scores = []
+        for batch in track(data_loader, description="Evaluating"):
+            input_ids, attention_mask = [b.to(self.device) for b in batch]
+            with torch.no_grad():
+                logits = self.model(
+                    input_ids=input_ids, attention_mask=attention_mask
+                ).logits
+            scaled_logits = logits / temperature
+            probabilities = F.softmax(scaled_logits, dim=-1)
+            positive_class_probabilities = (
+                probabilities[:, positive_label].cpu().numpy()
+            )
+            scores.extend(positive_class_probabilities)
+
+        return scores
+
+    def visualize_roc_curve(self, test_scores: list[float]):
+        plt.figure(figsize=(8, 6))
+        test_labels = [int(elt) for elt in self.test_dataset["label"]]
+        fpr, tpr, _ = roc_curve(test_labels, test_scores)
+        roc_auc = roc_auc_score(test_labels, test_scores)
+        plt.plot(
+            fpr,
+            tpr,
+            color="darkorange",
+            lw=2,
+            label=f"ROC curve (area = {roc_auc:.3f})",
+        )
+        plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.title("Receiver Operating Characteristic")
+        plt.legend(loc="lower right")
+        if self.streamlit_mode:
+            st.pyplot(plt)
+        else:
+            plt.show()
+
+    def evaluate_model(
+        self,
+        batch_size: int = 32,
+        positive_label: int = 2,
+        temperature: float = 3.0,
+        truncation: bool = True,
+        max_length: int = 512,
+    ):
+        test_scores = self.evaluate_batch(
+            self.test_dataset["text"],
+            batch_size=batch_size,
+            positive_label=positive_label,
+            temperature=temperature,
+            truncation=truncation,
+            max_length=max_length,
+        )
+        self.visualize_roc_curve(test_scores)
+        return test_scores

pyproject.toml

@@ -16,6 +16,9 @@ dependencies = [
     "transformers>=4.46.3",
     "torch>=2.5.1",
     "instructor>=1.7.0",
+    "matplotlib>=3.9.3",
+    "plotly>=5.24.1",
+    "scikit-learn>=1.5.2",
 ]
 
 [project.optional-dependencies]