fix: LlamaGuardFineTuner

application_pages/llama_guard_fine_tuning.py

@@ -11,8 +11,20 @@ def initialize_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
+    if "load_fine_tuner_button" not in st.session_state:
+        st.session_state.load_fine_tuner_button = False
+    if "is_fine_tuner_loaded" not in st.session_state:
+        st.session_state.is_fine_tuner_loaded = False
+    if "model_name" not in st.session_state:
+        st.session_state.model_name = ""
+    if "preview_dataset" not in st.session_state:
+        st.session_state.preview_dataset = False
+    if "evaluate_model" not in st.session_state:
+        st.session_state.evaluate_model = False
+    if "evaluation_batch_size" not in st.session_state:
+        st.session_state.evaluation_batch_size = None
+    if "evaluation_temperature" not in st.session_state:
+        st.session_state.evaluation_temperature = None
 
 
 initialize_session_state()
@@ -30,14 +42,39 @@ if st.session_state.dataset_address != "":
     )
     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,
+
+    model_name = st.sidebar.selectbox(
+        "Model Name",
+        ["meta-llama/Prompt-Guard-86M"],
+    )
+    st.session_state.model_name = model_name
+
+    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
+
+    load_fine_tuner_button = st.sidebar.button("Load Fine-Tuner")
+    st.session_state.load_fine_tuner_button = load_fine_tuner_button
+
+    if st.session_state.load_fine_tuner_button:
+        with st.status("Loading Fine-Tuner"):
+            st.session_state.llama_guard_fine_tuner.load_dataset(
+                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_model(
+                model_name=st.session_state.model_name
             )
-            st.session_state.llama_guard_fine_tuner.load_dataset(dataset_args)
-            st.session_state.llama_guard_fine_tuner.show_dataset_sample()
+            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,
+                )
+            st.session_state.is_fine_tuner_loaded = True

guardrails_genie/train/llama_guard.py

@@ -23,12 +23,14 @@ class LlamaGuardFineTuner:
         dataset = load_dataset(dataset_args.dataset_address)
         self.train_dataset = (
             dataset["train"]
-            if dataset_args.train_dataset_range > 0
+            if dataset_args.train_dataset_range <= 0
+            or dataset_args.train_dataset_range > len(dataset["train"])
             else dataset["train"].select(range(dataset_args.train_dataset_range))
         )
         self.test_dataset = (
             dataset["test"]
-            if dataset_args.test_dataset_range > 0
+            if dataset_args.test_dataset_range <= 0
+            or dataset_args.test_dataset_range > len(dataset["test"])
             else dataset["test"].select(range(dataset_args.test_dataset_range))
         )
 
@@ -69,7 +71,12 @@ class LlamaGuardFineTuner:
         data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
 
         scores = []
-        for batch in track(data_loader, description="Evaluating"):
+        progress_bar = (
+            st.progress(0, text="Evaluating") if self.streamlit_mode else None
+        )
+        for i, batch in track(
+            enumerate(data_loader), description="Evaluating", total=len(data_loader)
+        ):
             input_ids, attention_mask = [b.to(self.device) for b in batch]
             with torch.no_grad():
                 logits = self.model(
@@ -81,6 +88,12 @@ class LlamaGuardFineTuner:
                 probabilities[:, positive_label].cpu().numpy()
             )
             scores.extend(positive_class_probabilities)
+            if progress_bar:
+                progress_percentage = (i + 1) * 100 // len(data_loader)
+                progress_bar.progress(
+                    progress_percentage,
+                    text=f"Evaluating batch {i + 1}/{len(data_loader)}",
+                )
 
         return scores