Update app.py

app.py

@@ -1,28 +1,47 @@
-import gradio as gr
+"""
+app.py
+"""
+
+# Standard imports
+import json
 import os
-import openai
-import torch
 import sys
 import uuid
 from datetime import datetime
-import json
-
-from safetensors.torch import load_file
-from lionguard2 import LionGuard2, CATEGORIES
-from utils import get_embeddings
 
+# Third party imports
+import openai
+import gradio as gr
 import gspread
 from google.oauth2 import service_account
+from transformers import AutoModel
+
+# Local imports
+from utils import get_embeddings
+
+# --- Categories
+CATEGORIES = {
+    "binary": ["binary"],
+    "hateful": ["hateful_l1", "hateful_l2"],
+    "insults": ["insults"],
+    "sexual": [
+        "sexual_l1",
+        "sexual_l2",
+    ],
+    "physical_violence": ["physical_violence"],
+    "self_harm": ["self_harm_l1", "self_harm_l2"],
+    "all_other_misconduct": [
+        "all_other_misconduct_l1",
+        "all_other_misconduct_l2",
+    ],
+}
 
 # --- OpenAI Setup ---
 client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
 # --- Model Loading ---
 def load_lionguard2():
-    model = LionGuard2()
-    model.eval()
-    state_dict = load_file('LionGuard2.safetensors')
-    model.load_state_dict(state_dict)
+    model = AutoModel.from_pretrained("govtech/lionguard-2", trust_remote_code=True)
     return model
 
 model = load_lionguard2()