Update utils/risk_detector.py

utils/risk_detector.py

@@ -2,14 +2,30 @@
 
 from transformers import pipeline
 
+# ⚖️ Load zero-shot classification model
 classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli")
+
+# 🎯 Define risk-related labels
 labels = ["Indemnity", "Exclusivity", "Termination", "Jurisdiction", "Confidentiality"]
 
-def detect_risks(text):
-    if not text:
+def detect_risks(text, verbose=False):
+    """
+    Classify clauses into predefined legal risk categories.
+    If verbose=True, include detailed scores for each label.
+
+    Returns:
+        - List of (label, score) tuples (default)
+        - Or dict of full model output if verbose
+    """
+    if not text.strip():
         return []
 
     result = classifier(text[:1000], candidate_labels=labels, multi_label=True)
 
-    # Package into list of [label, score] for Gradio dataframe
+    if verbose:
+        return {
+            "sequence": result["sequence"],
+            "predictions": list(zip(result["labels"], result["scores"]))
+        }
+
     return list(zip(result["labels"], result["scores"]))