update: app

application_pages/train_classifier.py

@@ -1,3 +1,5 @@
+import os
+
 import streamlit as st
 from dotenv import load_dotenv
 
@@ -26,7 +28,12 @@ dataset_name = st.sidebar.text_input("Dataset Name", value="")
 st.session_state.dataset_name = dataset_name
 
 base_model_name = st.sidebar.selectbox(
-    "Base Model", options=["distilbert/distilbert-base-uncased", "roberta-base"]
+    "Base Model",
+    options=[
+        "distilbert/distilbert-base-uncased",
+        "FacebookAI/roberta-base",
+        "microsoft/deberta-v3-base",
+    ],
 )
 st.session_state.base_model_name = base_model_name
 
@@ -46,8 +53,9 @@ if st.session_state.should_start_training:
             f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
         )
         training_output = train_binary_classifier(
-            project_name="guardrails-genie",
-            entity_name="geekyrakshit",
+            project_name=os.getenv("WANDB_PROJECT_NAME"),
+            entity_name=os.getenv("WANDB_ENTITY_NAME"),
+            run_name=f"{st.session_state.base_model_name}-finetuned",
             dataset_repo=st.session_state.dataset_name,
             model_name=st.session_state.base_model_name,
             batch_size=st.session_state.batch_size,

guardrails_genie/train_classifier.py

@@ -1,4 +1,3 @@
-
 import evaluate
 import numpy as np
 import streamlit as st
@@ -39,6 +38,7 @@ class StreamlitProgressbarCallback(TrainerCallback):
 def train_binary_classifier(
     project_name: str,
     entity_name: str,
+    run_name: str,
     dataset_repo: str = "geekyrakshit/prompt-injection-dataset",
     model_name: str = "distilbert/distilbert-base-uncased",
     learning_rate: float = 2e-5,
@@ -47,7 +47,7 @@ def train_binary_classifier(
     weight_decay: float = 0.01,
     streamlit_mode: bool = False,
 ):
-    wandb.init(project=project_name, entity=entity_name)
+    wandb.init(project=project_name, entity=entity_name, name=run_name)
     dataset = load_dataset(dataset_repo)
     tokenizer = AutoTokenizer.from_pretrained(model_name)