Merge pull request #15 from soumik12345/fix/eval-app

README.md

@@ -18,11 +18,6 @@ source .venv/bin/activate
 ## Run the App
 
 ```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
 ```
 

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 
 intro_page = st.Page(
-    "application_pages/intro_page.py", title="Introduction", icon=":material/guardian:"
+    "application_pages/intro_page.py", title="Authenticate", icon=":material/guardian:"
 )
 chat_page = st.Page(
     "application_pages/chat_app.py",
@@ -13,23 +13,11 @@ evaluation_page = st.Page(
     title="Evaluation",
     icon=":material/monitoring:",
 )
-# train_classifier_page = st.Page(
-#     "application_pages/train_classifier.py",
-#     title="Train Classifier",
-#     icon=":material/fitness_center:",
-# )
-# llama_guard_fine_tuning_page = st.Page(
-#     "application_pages/llama_guard_fine_tuning.py",
-#     title="Fine-Tune LLama Guard",
-#     icon=":material/star:",
-# )
 page_navigation = st.navigation(
     [
         intro_page,
         chat_page,
         evaluation_page,
-        # train_classifier_page,
-        # llama_guard_fine_tuning_page,
     ]
 )
 st.set_page_config(page_title="Guardrails Genie", page_icon=":material/guardian:")

application_pages/chat_app.py

@@ -1,8 +1,6 @@
 import importlib
-import os
 
 import streamlit as st
-import weave
 from dotenv import load_dotenv
 
 from guardrails_genie.guardrails import GuardrailManager
@@ -11,8 +9,6 @@ from guardrails_genie.llm import OpenAIModel
 
 def initialize_session_state():
     load_dotenv()
-    weave.init(project_name=os.getenv("WEAVE_PROJECT"))
-
     if "guardrails" not in st.session_state:
         st.session_state.guardrails = []
     if "guardrail_names" not in st.session_state:
@@ -53,7 +49,6 @@ def initialize_guardrails():
                 [
                     "",
                     "ProtectAI/deberta-v3-base-prompt-injection-v2",
-                    "wandb://geekyrakshit/guardrails-genie/model-6rwqup9b:v3",
                 ],
             )
             if classifier_model_name != "":
@@ -120,63 +115,69 @@ def initialize_guardrails():
     )
 
 
-initialize_session_state()
-st.title(":material/robot: Guardrails Genie Playground")
-
-openai_model = st.sidebar.selectbox(
-    "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
-)
-chat_condition = openai_model != ""
+if st.session_state.is_authenticated:
+    initialize_session_state()
+    st.title(":material/robot: Guardrails Genie Playground")
 
-guardrails = []
-
-guardrail_names = st.sidebar.multiselect(
-    label="Select Guardrails",
-    options=[
-        cls_name
-        for cls_name, cls_obj in vars(
-            importlib.import_module("guardrails_genie.guardrails")
-        ).items()
-        if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
-    ],
-)
-st.session_state.guardrail_names = guardrail_names
+    openai_model = st.sidebar.selectbox(
+        "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
+    )
+    chat_condition = openai_model != ""
+
+    guardrails = []
+
+    guardrail_names = st.sidebar.multiselect(
+        label="Select Guardrails",
+        options=[
+            cls_name
+            for cls_name, cls_obj in vars(
+                importlib.import_module("guardrails_genie.guardrails")
+            ).items()
+            if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
+        ],
+    )
+    st.session_state.guardrail_names = guardrail_names
 
-if st.sidebar.button("Initialize Guardrails") and chat_condition:
-    st.session_state.initialize_guardrails = True
+    if st.sidebar.button("Initialize Guardrails") and chat_condition:
+        st.session_state.initialize_guardrails = True
 
-if st.session_state.initialize_guardrails:
-    with st.sidebar.status("Initializing Guardrails..."):
-        initialize_guardrails()
-        st.session_state.llm_model = OpenAIModel(model_name=openai_model)
+    if st.session_state.initialize_guardrails:
+        with st.sidebar.status("Initializing Guardrails..."):
+            initialize_guardrails()
+            st.session_state.llm_model = OpenAIModel(model_name=openai_model)
 
-    user_prompt = st.text_area("User Prompt", value="")
-    st.session_state.user_prompt = user_prompt
+        user_prompt = st.text_area("User Prompt", value="")
+        st.session_state.user_prompt = user_prompt
 
-    test_guardrails_button = st.button("Test Guardrails")
-    st.session_state.test_guardrails = test_guardrails_button
+        test_guardrails_button = st.button("Test Guardrails")
+        st.session_state.test_guardrails = test_guardrails_button
 
-    if st.session_state.test_guardrails:
-        with st.sidebar.status("Running Guardrails..."):
-            guardrails_response, call = st.session_state.guardrails_manager.guard.call(
-                st.session_state.guardrails_manager, prompt=st.session_state.user_prompt
-            )
+        if st.session_state.test_guardrails:
+            with st.sidebar.status("Running Guardrails..."):
+                guardrails_response, call = (
+                    st.session_state.guardrails_manager.guard.call(
+                        st.session_state.guardrails_manager,
+                        prompt=st.session_state.user_prompt,
+                    )
+                )
 
-        if guardrails_response["safe"]:
-            st.markdown(
-                f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
-            )
+            if guardrails_response["safe"]:
+                st.markdown(
+                    f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
+                )
 
-            with st.sidebar.status("Generating response from LLM..."):
-                response, call = st.session_state.llm_model.predict.call(
-                    st.session_state.llm_model,
-                    user_prompts=st.session_state.user_prompt,
+                with st.sidebar.status("Generating response from LLM..."):
+                    response, call = st.session_state.llm_model.predict.call(
+                        st.session_state.llm_model,
+                        user_prompts=st.session_state.user_prompt,
+                    )
+                st.markdown(
+                    response.choices[0].message.content
+                    + f"\n\n---\nExplore LLM generation trace on [Weave]({call.ui_url})"
                 )
-            st.markdown(
-                response.choices[0].message.content
-                + f"\n\n---\nExplore LLM generation trace on [Weave]({call.ui_url})"
-            )
-        else:
-            st.warning("Prompt is not safe!")
-            st.markdown(guardrails_response["summary"])
-            st.markdown(f"Explore prompt trace on [Weave]({call.ui_url})")
+            else:
+                st.warning("Prompt is not safe!")
+                st.markdown(guardrails_response["summary"])
+                st.markdown(f"Explore prompt trace on [Weave]({call.ui_url})")
+else:
+    st.warning("Please authenticate your WandB account to use this feature.")

application_pages/evaluation_app.py

@@ -1,10 +1,7 @@
 import asyncio
-import os
-import time
 from importlib import import_module
 
 import pandas as pd
-import rich
 import streamlit as st
 import weave
 from dotenv import load_dotenv
@@ -12,7 +9,6 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 from guardrails_genie.metrics import AccuracyMetric
-from guardrails_genie.utils import EvaluationCallManager
 
 
 def initialize_session_state():
@@ -20,48 +16,34 @@ def initialize_session_state():
     if "uploaded_file" not in st.session_state:
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
-        st.session_state.dataset_name = ""
+        st.session_state.dataset_name = None
     if "preview_in_app" not in st.session_state:
         st.session_state.preview_in_app = False
+    if "is_dataset_published" not in st.session_state:
+        st.session_state.is_dataset_published = False
+    if "publish_dataset_button" not in st.session_state:
+        st.session_state.publish_dataset_button = False
     if "dataset_ref" not in st.session_state:
         st.session_state.dataset_ref = None
-    if "dataset_previewed" not in st.session_state:
-        st.session_state.dataset_previewed = False
-    if "guardrail_names" not in st.session_state:
-        st.session_state.guardrail_names = []
     if "guardrails" not in st.session_state:
         st.session_state.guardrails = []
-    if "start_evaluation" not in st.session_state:
-        st.session_state.start_evaluation = False
-    if "evaluation_summary" not in st.session_state:
-        st.session_state.evaluation_summary = None
-    if "guardrail_manager" not in st.session_state:
-        st.session_state.guardrail_manager = None
+    if "guardrail_names" not in st.session_state:
+        st.session_state.guardrail_names = []
+    if "start_evaluations_button" not in st.session_state:
+        st.session_state.start_evaluations_button = False
     if "evaluation_name" not in st.session_state:
         st.session_state.evaluation_name = ""
-    if "show_result_table" not in st.session_state:
-        st.session_state.show_result_table = False
-    if "weave_client" not in st.session_state:
-        st.session_state.weave_client = weave.init(
-            project_name=os.getenv("WEAVE_PROJECT")
-        )
-    if "evaluation_call_manager" not in st.session_state:
-        st.session_state.evaluation_call_manager = None
-    if "call_id" not in st.session_state:
-        st.session_state.call_id = None
-    if "llama_guardrail_checkpoint" not in st.session_state:
-        st.session_state.llama_guardrail_checkpoint = None
-
-
-def initialize_guardrail():
-    guardrails = []
+
+
+def initialize_guardrails():
+    st.session_state.guardrails = []
     for guardrail_name in st.session_state.guardrail_names:
         if guardrail_name == "PromptInjectionSurveyGuardrail":
             survey_guardrail_model = st.sidebar.selectbox(
                 "Survey Guardrail LLM", ["", "gpt-4o-mini", "gpt-4o"]
             )
             if survey_guardrail_model:
-                guardrails.append(
+                st.session_state.guardrails.append(
                     getattr(
                         import_module("guardrails_genie.guardrails"),
                         guardrail_name,
@@ -73,29 +55,60 @@ def initialize_guardrail():
                 [
                     "",
                     "ProtectAI/deberta-v3-base-prompt-injection-v2",
-                    "wandb://geekyrakshit/guardrails-genie/model-6rwqup9b:v3",
                 ],
             )
-            if classifier_model_name:
+            if classifier_model_name != "":
                 st.session_state.guardrails.append(
                     getattr(
                         import_module("guardrails_genie.guardrails"),
                         guardrail_name,
                     )(model_name=classifier_model_name)
                 )
+        elif guardrail_name == "PresidioEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "RegexEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "TransformersEntityRecognitionGuardrail":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
+        elif guardrail_name == "RestrictedTermsJudge":
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(should_anonymize=True)
+            )
         elif guardrail_name == "PromptInjectionLlamaGuardrail":
-            llama_guardrail_checkpoint = st.sidebar.text_input(
-                "Llama Guardrail Checkpoint",
-                value=None,
+            llama_guard_checkpoint_name = st.sidebar.text_input(
+                "Checkpoint Name", value=""
             )
-            st.session_state.llama_guardrail_checkpoint = llama_guardrail_checkpoint
-            if st.session_state.llama_guardrail_checkpoint is not None:
-                st.session_state.guardrails.append(
-                    getattr(
-                        import_module("guardrails_genie.guardrails"),
-                        guardrail_name,
-                    )(checkpoint=st.session_state.llama_guardrail_checkpoint)
+            st.session_state.llama_guard_checkpoint_name = llama_guard_checkpoint_name
+            st.session_state.guardrails.append(
+                getattr(
+                    import_module("guardrails_genie.guardrails"),
+                    guardrail_name,
+                )(
+                    checkpoint=(
+                        None
+                        if st.session_state.llama_guard_checkpoint_name == ""
+                        else st.session_state.llama_guard_checkpoint_name
+                    )
                 )
+            )
         else:
             st.session_state.guardrails.append(
                 getattr(
@@ -103,64 +116,79 @@ def initialize_guardrail():
                     guardrail_name,
                 )()
             )
-    st.session_state.guardrails = guardrails
-    st.session_state.guardrail_manager = GuardrailManager(guardrails=guardrails)
-
-
-initialize_session_state()
-st.title(":material/monitoring: Evaluation")
-
-uploaded_file = st.sidebar.file_uploader(
-    "Upload the evaluation dataset as a CSV file", type="csv"
-)
-st.session_state.uploaded_file = uploaded_file
-dataset_name = st.sidebar.text_input("Evaluation dataset name", value="")
-st.session_state.dataset_name = dataset_name
-preview_in_app = st.sidebar.toggle("Preview in app", value=False)
-st.session_state.preview_in_app = preview_in_app
-
-if st.session_state.uploaded_file is not None and st.session_state.dataset_name != "":
-    with st.expander("Evaluation Dataset Preview", expanded=True):
-        dataframe = pd.read_csv(st.session_state.uploaded_file)
-        data_list = dataframe.to_dict(orient="records")
-
-        dataset = weave.Dataset(name=st.session_state.dataset_name, rows=data_list)
-        st.session_state.dataset_ref = weave.publish(dataset)
-
-        entity = st.session_state.dataset_ref.entity
-        project = st.session_state.dataset_ref.project
-        dataset_name = st.session_state.dataset_name
-        digest = st.session_state.dataset_ref._digest
-        st.markdown(
-            f"Dataset published to [**Weave**](https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest})"
-        )
-
-        if preview_in_app:
-            st.dataframe(dataframe)
-
-        st.session_state.dataset_previewed = True
-
-if st.session_state.dataset_previewed:
-    guardrail_names = st.sidebar.multiselect(
-        "Select Guardrails",
-        options=[
-            cls_name
-            for cls_name, cls_obj in vars(
-                import_module("guardrails_genie.guardrails")
-            ).items()
-            if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
-        ],
+    st.session_state.guardrails_manager = GuardrailManager(
+        guardrails=st.session_state.guardrails
+    )
+
+
+if st.session_state.is_authenticated:
+    initialize_session_state()
+    st.title(":material/monitoring: Evaluation")
+
+    uploaded_file = st.sidebar.file_uploader(
+        "Upload the evaluation dataset as a CSV file", type="csv"
     )
-    st.session_state.guardrail_names = guardrail_names
-
-    if st.session_state.guardrail_names != []:
-        initialize_guardrail()
-        evaluation_name = st.sidebar.text_input("Evaluation name", value="")
-        st.session_state.evaluation_name = evaluation_name
-        if st.session_state.guardrail_manager is not None:
-            if st.sidebar.button("Start Evaluation"):
-                st.session_state.start_evaluation = True
-            if st.session_state.start_evaluation:
+    st.session_state.uploaded_file = uploaded_file
+
+    if st.session_state.uploaded_file is not None:
+        dataset_name = st.sidebar.text_input("Evaluation dataset name", value=None)
+        st.session_state.dataset_name = dataset_name
+        preview_in_app = st.sidebar.toggle("Preview in app", value=False)
+        st.session_state.preview_in_app = preview_in_app
+        publish_dataset_button = st.sidebar.button("Publish dataset")
+        st.session_state.publish_dataset_button = publish_dataset_button
+
+        if st.session_state.publish_dataset_button and (
+            st.session_state.dataset_name is not None
+            and st.session_state.dataset_name != ""
+        ):
+
+            with st.expander("Evaluation Dataset Preview", expanded=True):
+                dataframe = pd.read_csv(st.session_state.uploaded_file)
+                data_list = dataframe.to_dict(orient="records")
+
+                dataset = weave.Dataset(
+                    name=st.session_state.dataset_name, rows=data_list
+                )
+                st.session_state.dataset_ref = weave.publish(dataset)
+
+                entity = st.session_state.dataset_ref.entity
+                project = st.session_state.dataset_ref.project
+                dataset_name = st.session_state.dataset_name
+                digest = st.session_state.dataset_ref._digest
+                dataset_url = f"https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest}"
+                st.markdown(f"Dataset published to [**Weave**]({dataset_url})")
+
+                if preview_in_app:
+                    st.dataframe(dataframe.head(20))
+                    if len(dataframe) > 20:
+                        st.markdown(
+                            f"⚠️ Dataset is too large to preview in app, please explore in the [**Weave UI**]({dataset_url})"
+                        )
+
+            st.session_state.is_dataset_published = True
+
+        if st.session_state.is_dataset_published:
+            guardrail_names = st.sidebar.multiselect(
+                "Select Guardrails",
+                options=[
+                    cls_name
+                    for cls_name, cls_obj in vars(
+                        import_module("guardrails_genie.guardrails")
+                    ).items()
+                    if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
+                ],
+            )
+            st.session_state.guardrail_names = guardrail_names
+
+            initialize_guardrails()
+            evaluation_name = st.sidebar.text_input("Evaluation Name", value="")
+            st.session_state.evaluation_name = evaluation_name
+
+            start_evaluations_button = st.sidebar.button("Start Evaluations")
+            st.session_state.start_evaluations_button = start_evaluations_button
+            if st.session_state.start_evaluations_button:
+                # st.write(len(st.session_state.guardrails))
                 evaluation = weave.Evaluation(
                     dataset=st.session_state.dataset_ref,
                     scorers=[AccuracyMetric()],
@@ -170,10 +198,14 @@ if st.session_state.dataset_previewed:
                     evaluation_summary, call = asyncio.run(
                         evaluation.evaluate.call(
                             evaluation,
-                            st.session_state.guardrail_manager,
+                            GuardrailManager(guardrails=st.session_state.guardrails),
                             __weave={
-                                "display_name": "Evaluation.evaluate:"
-                                + st.session_state.evaluation_name
+                                "display_name": (
+                                    "Evaluation.evaluate"
+                                    if st.session_state.evaluation_name == ""
+                                    else "Evaluation.evaluate:"
+                                    + st.session_state.evaluation_name
+                                )
                             },
                         )
                     )
@@ -187,37 +219,8 @@ if st.session_state.dataset_previewed:
                         x="Metric",
                         y="Score",
                     )
-                    st.session_state.evaluation_summary = evaluation_summary
-                    st.session_state.call_id = call.id
-                    st.session_state.start_evaluation = False
-
-                    if not st.session_state.start_evaluation:
-                        time.sleep(5)
-                        st.session_state.evaluation_call_manager = (
-                            EvaluationCallManager(
-                                entity="geekyrakshit",
-                                project="guardrails-genie",
-                                call_id=st.session_state.call_id,
-                            )
-                        )
-                        for guardrail_name in st.session_state.guardrail_names:
-                            st.session_state.evaluation_call_manager.call_list.append(
-                                {
-                                    "guardrail_name": guardrail_name,
-                                    "calls": st.session_state.evaluation_call_manager.collect_guardrail_guard_calls_from_eval(),
-                                }
-                            )
-                            rich.print(
-                                st.session_state.evaluation_call_manager.call_list
-                            )
-                        st.dataframe(
-                            st.session_state.evaluation_call_manager.render_calls_to_streamlit()
-                        )
-                        if st.session_state.evaluation_call_manager.show_warning_in_app:
-                            st.warning(
-                                f"Only {st.session_state.evaluation_call_manager.max_count} calls can be shown in the app."
-                            )
-                        st.markdown(
-                            f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
-                        )
-                    st.session_state.evaluation_call_manager = None
+                    st.markdown(
+                        f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
+                    )
+else:
+    st.warning("Please authenticate your WandB account to use this feature.")

application_pages/intro_page.py

@@ -1,7 +1,54 @@
+import os
+
 import streamlit as st
 
+import wandb
+
+
+def initialize_session_state():
+    if "weave_project_name" not in st.session_state:
+        st.session_state.weave_project_name = "guardrails-genie"
+    if "weave_entity_name" not in st.session_state:
+        st.session_state.weave_entity_name = ""
+    if "wandb_api_key" not in st.session_state:
+        st.session_state.wandb_api_key = ""
+    if "authenticate_button" not in st.session_state:
+        st.session_state.authenticate_button = False
+    if "is_authenticated" not in st.session_state:
+        st.session_state.is_authenticated = False
+
+
+initialize_session_state()
 st.title("🧞‍♂️ Guardrails Genie")
 
 st.write(
     "Guardrails-Genie is a tool that helps you implement guardrails in your LLM applications."
 )
+
+st.sidebar.markdown(
+    "Get your Wandb API key from [https://wandb.ai/authorize](https://wandb.ai/authorize)"
+)
+weave_entity_name = st.sidebar.text_input(
+    "Weave Entity Name", value=st.session_state.weave_entity_name
+)
+st.session_state.weave_entity_name = weave_entity_name
+weave_project_name = st.sidebar.text_input(
+    "Weave Project Name", value=st.session_state.weave_project_name
+)
+st.session_state.weave_project_name = weave_project_name
+wandb_api_key = st.sidebar.text_input("Wandb API Key", value="", type="password")
+st.session_state.wandb_api_key = wandb_api_key
+openai_api_key = st.sidebar.text_input("OpenAI API Key", value="", type="password")
+os.environ["OPENAI_API_KEY"] = openai_api_key
+authenticate_button = st.sidebar.button("Authenticate")
+st.session_state.authenticate_button = authenticate_button
+
+if authenticate_button and (
+    st.session_state.wandb_api_key != "" and st.session_state.weave_project_name != ""
+):
+    is_wandb_logged_in = wandb.login(key=st.session_state.wandb_api_key, relogin=True)
+    if is_wandb_logged_in:
+        st.session_state.is_authenticated = True
+        st.success("Logged in to Wandb")
+    else:
+        st.error("Failed to log in to Wandb")

guardrails_genie/train/llama_guard.py

@@ -3,12 +3,13 @@ import shutil
 from glob import glob
 from typing import Optional
 
+# import torch.optim as optim
+import bitsandbytes.optim as optim
 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
 from datasets import load_dataset
 from pydantic import BaseModel
 from rich.progress import track
@@ -335,8 +336,8 @@ class LlamaGuardFineTuner:
 
     def train(
         self,
-        batch_size: int = 32,
-        lr: float = 5e-6,
+        batch_size: int = 16,
+        starting_lr: float = 1e-7,
         num_classes: int = 2,
         log_interval: int = 1,
         save_interval: int = 50,
@@ -358,7 +359,7 @@ class LlamaGuardFineTuner:
 
         Args:
             batch_size (int, optional): The number of samples per batch during training.
-            lr (float, optional): The learning rate for the optimizer.
+            starting_lr (float, optional): The starting learning rate for the optimizer.
             num_classes (int, optional): The number of output classes for the classifier.
             log_interval (int, optional): The interval (in batches) at which to log the loss.
             save_interval (int, optional): The interval (in batches) at which to save model checkpoints.
@@ -377,7 +378,7 @@ class LlamaGuardFineTuner:
         wandb.config.dataset_args = self.dataset_args.model_dump()
         wandb.config.model_name = self.model_name
         wandb.config.batch_size = batch_size
-        wandb.config.lr = lr
+        wandb.config.starting_lr = starting_lr
         wandb.config.num_classes = num_classes
         wandb.config.log_interval = log_interval
         wandb.config.save_interval = save_interval
@@ -387,7 +388,16 @@ class LlamaGuardFineTuner:
         self.model.num_labels = num_classes
         self.model = self.model.to(self.device)
         self.model.train()
-        optimizer = optim.AdamW(self.model.parameters(), lr=lr)
+        # optimizer = optim.AdamW(self.model.parameters(), lr=starting_lr)
+        optimizer = optim.Lion(
+            self.model.parameters(), lr=starting_lr, weight_decay=0.01
+        )
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=starting_lr,
+            steps_per_epoch=len(self.train_dataset) // batch_size + 1,
+            epochs=1,
+        )
         data_loader = DataLoader(
             self.train_dataset,
             batch_size=batch_size,
@@ -405,9 +415,14 @@ class LlamaGuardFineTuner:
             loss = outputs.loss
             optimizer.zero_grad()
             loss.backward()
+
+            # torch.nn.utils.clip_grad_norm_(self.model.parameters(), gradient_clipping)
+
             optimizer.step()
+            scheduler.step()
             if (i + 1) % log_interval == 0:
                 wandb.log({"loss": loss.item()}, step=i + 1)
+                wandb.log({"learning_rate": scheduler.get_last_lr()[0]}, step=i + 1)
             if progress_bar:
                 progress_percentage = (i + 1) * 100 // len(data_loader)
                 progress_bar.progress(