Update app.py

app.py

@@ -179,38 +179,39 @@ def predict_batch(contents, policies):
     with torch.inference_mode():
         outputs = model(input_ids)
 
-        # Get logits for the last tokens
-        logits = outputs.logits[:, -1, :]
+    # Get logits for the last tokens
+    logits = outputs.logits[:, -1, :]
 
-        # Get token IDs for "0" and "1"
-        token_id_0 = tokenizer.encode("0", add_special_tokens=False)[0]
-        token_id_1 = tokenizer.encode("1", add_special_tokens=False)[0]
+    # Get token IDs for "0" and "1"
+    token_id_0 = tokenizer.encode("0", add_special_tokens=False)[0]
+    token_id_1 = tokenizer.encode("1", add_special_tokens=False)[0]
 
-        # Extract logits for "0" and "1"
-        binary_logits = logits[:, [token_id_0, token_id_1]]
+    # Extract logits for "0" and "1"
+    binary_logits = logits[:, [token_id_0, token_id_1]]
 
-        # Apply softmax to get probabilities for these two tokens
-        probabilities = F.softmax(binary_logits, dim=-1)
+    # Apply softmax to get probabilities for these two tokens
+    probabilities = F.softmax(binary_logits, dim=-1)
 
-        probs_0 = probabilities[:, 0].cpu().numpy()
-        probs_1 = probabilities[:, 1].cpu().numpy()
+    probs_0 = probabilities[:, 0].cpu().numpy()
+    probs_1 = probabilities[:, 1].cpu().numpy()
 
-        results = []
-        for prob_0, prob_1 in zip(probs_0, probs_1):
-            if prob_1 > prob_0:
-                output = f'VIOLATING\n(P: {prob_1:.2f})'
-            else:
-                output = f'NON-Violating\n(P: {prob_0:.2f})'
-            results.append(output)
-        print(results)
-        return results
+    results = []
+    for prob_0, prob_1 in zip(probs_0, probs_1):
+        if prob_1 > prob_0:
+            output = f'VIOLATING\n(P: {prob_1:.2f})'
+        else:
+            output = f'NON-Violating\n(P: {prob_0:.2f})'
+        results.append(output)
+    print(results)
+    #return results
+    return "\n\n".join(results)
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=predict_batch,
     inputs=[gr.Textbox(label="Content", lines=2, value=DEFAULT_CONTENT), 
             gr.Textbox(label="Policy", lines=10, value=DEFAULT_POLICY)], 
-    outputs=[gr.Textbox(label="Results")],
+    outputs=gr.Textbox(label="Results"),
     batch=True,
     max_batch_size=8,
     title="CoPE Dev (Unstable)",