Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import os
 
 import torch
+import torch.nn.functional as F
 from peft import PeftConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
@@ -97,13 +98,34 @@ def predict(content, policy):
 
     with torch.inference_mode():
         outputs = model(input_ids)
-        logits = outputs.logits[:, -1, :]  # Get logits for the last token
+
+        # Get logits for the last token
+        logits = outputs.logits[:, -1, :]  
+
+        # Apply softmax to get probabilities
+        probabilities = F.softmax(logits, dim=-1)
+        
+        # Get the predicted token ID
         predicted_token_id = torch.argmax(logits, dim=-1).item()
-    decoded_output = tokenizer.decode([predicted_token_id])
-    if decoded_output == '0':
-        return f'NON-Violating ({decoded_output})'
+        
+        # Decode the predicted token
+        decoded_output = tokenizer.decode([predicted_token_id])
+        
+        # Get the probability of the predicted token
+        predicted_prob = probabilities[0, predicted_token_id].item()
+        
+        # Function to get probability for a specific token
+        def get_token_probability(token):
+            token_id = tokenizer.encode(token, add_special_tokens=False)[0]
+            return probabilities[0, token_id].item()
+
+        predicted_prob_0 = get_token_probability('0')
+        predicted_prob_1 = get_token_probability('1')
+        
+    if decoded_output == '1':
+        return f'VIOLATING\n(P: {predicted_prob_1:.2f})'
     else:
-        return f'VIOLATING ({decoded_output})'
+        return f'NON-Violating\n(P: {predicted_prob_0:.2f})'
 
 with gr.Blocks() as iface:
     gr.Markdown("# CoPE Alpha Preview")
@@ -127,7 +149,7 @@ with gr.Blocks() as iface:
     2. Specify your policy in the "Policy" box.
     3. Click "Submit" to see the results.
 
-    **Note**: Inference times are **slow** (2-5 seconds) since this is built on dev infra and not yet optimized for live systems. Please be patient!
+    **Note**: Inference times are **slow** (1-2 seconds) since this is built on dev infra and not yet optimized for live systems. Please be patient!
     
     ## More Info