proverka obuchenya

.ipynb_checkpoints/proverkabert-checkpoint.py

@@ -183,9 +183,11 @@ def classify_prompt(prompt, tokenizer, model):
         "safety": safety_label,
         "attack_type": attack_label,
         "safety_confidence": round(probs_safety[0, pred_safety].item(), 4),
-        "attack_confidence": round(probs_attack[0, pred_attack].item(), 4) if safety_label == "unsafe" else 0.0
+        "attack_confidence": round(probs_attack[0, pred_attack].item(), 4) if safety_label == "unsafe" else 0.0,
+        "safety_probs": probs_safety.tolist(),
+        "attack_probs": probs_attack.tolist()
     }
-
+    
 def main():
     MODEL_PATH = "./fine-tuned-bert-lora-multi-task"
     

proverkabert.py

@@ -183,9 +183,11 @@ def classify_prompt(prompt, tokenizer, model):
         "safety": safety_label,
         "attack_type": attack_label,
         "safety_confidence": round(probs_safety[0, pred_safety].item(), 4),
-        "attack_confidence": round(probs_attack[0, pred_attack].item(), 4) if safety_label == "unsafe" else 0.0
+        "attack_confidence": round(probs_attack[0, pred_attack].item(), 4) if safety_label == "unsafe" else 0.0,
+        "safety_probs": probs_safety.tolist(),
+        "attack_probs": probs_attack.tolist()
     }
-
+    
 def main():
     MODEL_PATH = "./fine-tuned-bert-lora-multi-task"