add: LlamaGuardFineTuner.train

guardrails_genie/train/llama_guard.py

@@ -1,11 +1,18 @@
+import os
+
 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
+import wandb
 from datasets import load_dataset
 from pydantic import BaseModel
 from rich.progress import track
+from safetensors.torch import save_model
 from sklearn.metrics import roc_auc_score, roc_curve
+from torch.utils.data import DataLoader
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 
@@ -16,7 +23,11 @@ class DatasetArgs(BaseModel):
 
 
 class LlamaGuardFineTuner:
-    def __init__(self, streamlit_mode: bool = False):
+    def __init__(
+        self, wandb_project: str, wandb_entity: str, streamlit_mode: bool = False
+    ):
+        self.wandb_project = wandb_project
+        self.wandb_entity = wandb_entity
         self.streamlit_mode = streamlit_mode
 
     def load_dataset(self, dataset_args: DatasetArgs):
@@ -36,6 +47,7 @@ class LlamaGuardFineTuner:
 
     def load_model(self, model_name: str = "meta-llama/Prompt-Guard-86M"):
         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
@@ -101,7 +113,6 @@ class LlamaGuardFineTuner:
         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)
-
         fig = go.Figure()
         fig.add_trace(
             go.Scatter(
@@ -121,7 +132,6 @@ class LlamaGuardFineTuner:
                 line=dict(color="navy", width=2, dash="dash"),
             )
         )
-
         fig.update_layout(
             title="Receiver Operating Characteristic",
             xaxis_title="False Positive Rate",
@@ -130,7 +140,6 @@ class LlamaGuardFineTuner:
             yaxis=dict(range=[0.0, 1.05]),
             legend=dict(x=0.8, y=0.2),
         )
-
         if self.streamlit_mode:
             st.plotly_chart(fig)
         else:
@@ -140,10 +149,7 @@ class LlamaGuardFineTuner:
         test_labels = [int(elt) for elt in self.test_dataset["label"]]
         positive_scores = [scores[i] for i in range(500) if test_labels[i] == 1]
         negative_scores = [scores[i] for i in range(500) if test_labels[i] == 0]
-
         fig = go.Figure()
-
-        # Plotting positive scores
         fig.add_trace(
             go.Histogram(
                 x=positive_scores,
@@ -153,8 +159,6 @@ class LlamaGuardFineTuner:
                 opacity=0.75,
             )
         )
-
-        # Plotting negative scores
         fig.add_trace(
             go.Histogram(
                 x=negative_scores,
@@ -164,8 +168,6 @@ class LlamaGuardFineTuner:
                 opacity=0.75,
             )
         )
-
-        # Updating layout
         fig.update_layout(
             title="Score Distribution for Positive and Negative Examples",
             xaxis_title="Score",
@@ -173,8 +175,6 @@ class LlamaGuardFineTuner:
             barmode="overlay",
             legend_title="Scores",
         )
-
-        # Display the plot
         if self.streamlit_mode:
             st.plotly_chart(fig)
         else:
@@ -199,3 +199,53 @@ class LlamaGuardFineTuner:
         self.visualize_roc_curve(test_scores)
         self.visualize_score_distribution(test_scores)
         return test_scores
+
+    def collate_fn(self, batch):
+        texts = [item["text"] for item in batch]
+        labels = torch.tensor([int(item["label"]) for item in batch])
+        encodings = self.tokenizer(
+            texts, padding=True, truncation=True, max_length=512, return_tensors="pt"
+        )
+        return encodings.input_ids, encodings.attention_mask, labels
+
+    def train(self, batch_size: int = 32, lr: float = 5e-6, num_classes: int = 2):
+        wandb.init(
+            project=self.wandb_project,
+            entity=self.wandb_entity,
+            name=f"{self.model_name}-{self.dataset_name}",
+        )
+        self.model.classifier = nn.Linear(
+            self.model.classifier.in_features, num_classes
+        )
+        self.model.num_labels = num_classes
+        self.model.train()
+        optimizer = optim.AdamW(self.model.parameters(), lr=lr)
+        data_loader = DataLoader(
+            self.train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            collate_fn=self.collate_fn,
+        )
+        progress_bar = st.progress(0, text="Training") if self.streamlit_mode else None
+        for i, batch in track(
+            enumerate(data_loader), description="Training", total=len(data_loader)
+        ):
+            input_ids, attention_mask, labels = [x.to(self.device) for x in batch]
+            outputs = self.model(
+                input_ids=input_ids, attention_mask=attention_mask, labels=labels
+            )
+            loss = outputs.loss
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            wandb.log({"loss": loss.item()})
+            if progress_bar:
+                progress_percentage = (i + 1) * 100 // len(data_loader)
+                progress_bar.progress(
+                    progress_percentage,
+                    text=f"Training batch {i + 1}/{len(data_loader)}, Loss: {loss.item()}",
+                )
+        save_model(self.model, f"{self.model_name}-{self.dataset_name}.safetensors")
+        wandb.log_model(f"{self.model_name}-{self.dataset_name}.safetensors")
+        wandb.finish()
+        os.remove(f"{self.model_name}-{self.dataset_name}.safetensors")