update: app

README.md

@@ -20,6 +20,8 @@ source .venv/bin/activate
 ```bash
 export OPENAI_API_KEY="YOUR_OPENAI_API_KEY"
 export WEAVE_PROJECT="YOUR_WEAVE_PROJECT"
+export WANDB_PROJECT_NAME="YOUR_WANDB_PROJECT_NAME"
+export WANDB_ENTITY_NAME="YOUR_WANDB_ENTITY_NAME"
 export WANDB_LOG_MODEL="checkpoint"
 streamlit run app.py
 ```

application_pages/train_classifier.py

@@ -3,7 +3,6 @@ import os
 import streamlit as st
 from dotenv import load_dotenv
 
-import wandb
 from guardrails_genie.train_classifier import train_binary_classifier
 
 
@@ -49,9 +48,6 @@ st.session_state.should_start_training = (
 
 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=os.getenv("WANDB_PROJECT_NAME"),
             entity_name=os.getenv("WANDB_ENTITY_NAME"),

guardrails_genie/train_classifier.py

@@ -48,6 +48,10 @@ def train_binary_classifier(
     streamlit_mode: bool = False,
 ):
     wandb.init(project=project_name, entity=entity_name, name=run_name)
+    if streamlit_mode:
+        st.markdown(
+            f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
+        )
     dataset = load_dataset(dataset_repo)
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -97,6 +101,10 @@ def train_binary_classifier(
         compute_metrics=compute_metrics,
         callbacks=[StreamlitProgressbarCallback()] if streamlit_mode else [],
     )
-    training_output = trainer.train()
+    try:
+        training_output = trainer.train()
+    except Exception as e:
+        wandb.finish()
+        raise e
     wandb.finish()
     return training_output