add: integration of training with app

app.py

@@ -13,6 +13,13 @@ evaluation_page = st.Page(
     title="Evaluation",
     icon=":material/monitoring:",
 )
-page_navigation = st.navigation([intro_page, chat_page, evaluation_page])
+train_classifier_page = st.Page(
+    "application_pages/train_classifier.py",
+    title="Train Classifier",
+    icon=":material/fitness_center:",
+)
+page_navigation = st.navigation(
+    [intro_page, chat_page, evaluation_page, train_classifier_page]
+)
 st.set_page_config(page_title="Guardrails Genie", page_icon=":material/guardian:")
 page_navigation.run()

application_pages/train_classifier.py

@@ -0,0 +1,57 @@
+import streamlit as st
+from dotenv import load_dotenv
+
+import wandb
+from guardrails_genie.train_classifier import train_binary_classifier
+
+
+def initialize_session_state():
+    load_dotenv()
+    if "dataset_name" not in st.session_state:
+        st.session_state.dataset_name = None
+    if "base_model_name" not in st.session_state:
+        st.session_state.base_model_name = None
+    if "batch_size" not in st.session_state:
+        st.session_state.batch_size = 16
+    if "should_start_training" not in st.session_state:
+        st.session_state.should_start_training = False
+    if "training_output" not in st.session_state:
+        st.session_state.training_output = None
+
+
+initialize_session_state()
+st.title(":material/fitness_center: Train Classifier")
+
+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"]
+)
+st.session_state.base_model_name = base_model_name
+
+batch_size = st.sidebar.slider(
+    "Batch Size", min_value=4, max_value=256, value=16, step=4
+)
+st.session_state.batch_size = batch_size
+
+train_button = st.sidebar.button("Train")
+st.session_state.should_start_training = (
+    train_button and st.session_state.dataset_name and st.session_state.base_model_name
+)
+
+if st.session_state.should_start_training:
+    with st.expander("Training", expanded=True):
+        st.markdown(
+            f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
+        )
+        training_output = train_binary_classifier(
+            project_name="guardrails-genie",
+            entity_name="geekyrakshit",
+            dataset_repo=st.session_state.dataset_name,
+            model_name=st.session_state.base_model_name,
+            batch_size=st.session_state.batch_size,
+            streamlit_mode=True,
+        )
+        st.session_state.training_output = training_output
+        st.write(training_output)

guardrails_genie/train_classifier.py

@@ -1,14 +1,39 @@
+
 import evaluate
 import numpy as np
-import wandb
+import streamlit as st
 from datasets import load_dataset
 from transformers import (
     AutoModelForSequenceClassification,
     AutoTokenizer,
     DataCollatorWithPadding,
     Trainer,
+    TrainerCallback,
     TrainingArguments,
 )
+from transformers.trainer_callback import TrainerControl, TrainerState
+
+import wandb
+
+
+class StreamlitProgressbarCallback(TrainerCallback):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.progress_bar = st.progress(0, text="Training")
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        super().on_step_begin(args, state, control, **kwargs)
+        self.progress_bar.progress(
+            (state.global_step * 100 // state.max_steps) + 1,
+            text=f"Training {state.global_step} / {state.max_steps}",
+        )
 
 
 def train_binary_classifier(
@@ -20,6 +45,7 @@ def train_binary_classifier(
     batch_size: int = 16,
     num_epochs: int = 2,
     weight_decay: float = 0.01,
+    streamlit_mode: bool = False,
 ):
     wandb.init(project=project_name, entity=entity_name)
     dataset = load_dataset(dataset_repo)
@@ -69,5 +95,8 @@ def train_binary_classifier(
         processing_class=tokenizer,
         data_collator=data_collator,
         compute_metrics=compute_metrics,
+        callbacks=[StreamlitProgressbarCallback()] if streamlit_mode else [],
     )
-    trainer.train()
+    training_output = trainer.train()
+    wandb.finish()
+    return training_output