test: Testing system prompt

app.py

@@ -15,26 +15,20 @@ 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"
@@ -44,9 +38,15 @@ def check_safety(text):
 def inference(prompt: str, max_tokens: int = 256) -> str:
     model.to("cuda")  # Move to GPU inside decorated function
 
-    messages = [
-        {"role" : "user", "content" : prompt}
-    ]
+    SYSTEM_PROMPT = """You are having a casual conversation with a user in Singapore.
+     Keep responses helpful and friendly. Avoid sensitive topics like politics, religion, or race.
+     If asked about harmful activities, politely decline."""
+    
+     messages = [
+         {"role": "system", "content": SYSTEM_PROMPT},
+         {"role": "user", "content": prompt}
+     ]
+
     text = tokenizer.apply_chat_template(
         messages,
         tokenize = False,