feat: Added LionGuard v1 for safety detector

app.py

@@ -2,8 +2,10 @@ import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import torch
-
-# torch.manual_seed(42)
+import onnxruntime as ort
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import hf_hub_download
 
 MODEL_ID = "yuhueng/qwen3-4b-singlish-base"  # replace with your model
 
@@ -13,6 +15,31 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16,
 )
 
+# --- 1. Configuration ---
+REPO_ID = "govtech/lionguard-v1"
+EMBEDDING_MODEL = "BAAI/bge-large-en-v1.5" 
+FILENAME = "models/lionguard-binary.onnx"
+
+# --- 2. Load Models ---
+embedder = SentenceTransformer(EMBEDDING_MODEL)
+
+model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+session = ort.InferenceSession(model_path)
+
+# --- 3. The Inference Logic ---
+def check_safety(text):
+    # Generate embedding (Normalize is important for BGE models)
+    embedding = embedder.encode([text], normalize_embeddings=True)
+
+    # Prepare input for ONNX
+    input_name = session.get_inputs()[0].name
+    
+    # Run prediction
+    pred = session.run(None, {input_name: embedding.astype(np.float32)})[0]
+    
+    return "Unsafe" if pred[0] == 1 else "Safe"
+
+
 @spaces.GPU(duration=120)
 def inference(prompt: str, max_tokens: int = 256) -> str:
     model.to("cuda")  # Move to GPU inside decorated function
@@ -40,7 +67,8 @@ def inference(prompt: str, max_tokens: int = 256) -> str:
         outputs[0][inputs["input_ids"].shape[1]:],
         skip_special_tokens=True
     )
-    return response
+    safety = check_safety(response)
+    return response, safety
     
     # # Use TextIteratorStreamer instead of TextStreamer
     # streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)