Wire trained Legal-BERT model (76% F1) into live demo — replaces regex

app.py

@@ -1,33 +1,74 @@
 """
 ClauseGuard — AI Fine Print Scanner
+Uses Legal-BERT fine-tuned on CLAUDETTE/LexGLUE unfair_tos (8 categories).
 """
 
 import gradio as gr
 import re
+import numpy as np
+
+# ─── Load ML model ───
+MODEL_ID = "gaurv007/clauseguard-legal-bert"
+ml_pipeline = None
+
+try:
+    from transformers import pipeline
+    ml_pipeline = pipeline("text-classification", model=MODEL_ID, top_k=None, device=-1)
+    print(f"Loaded model: {MODEL_ID}")
+except Exception as e:
+    print(f"Model load failed ({e}), using regex fallback")
+
+# ─── Label metadata ───
+LABELS = {
+    "Limitation of liability": ("HIGH", "Company avoids responsibility for damages or losses."),
+    "Unilateral termination": ("HIGH", "They can close your account without reason."),
+    "Unilateral change": ("MEDIUM", "Terms can change without your consent."),
+    "Content removal": ("MEDIUM", "Your content can be deleted without notice."),
+    "Contract by using": ("LOW", "You agree just by visiting or using the site."),
+    "Choice of law": ("MEDIUM", "Foreign law applies instead of your local protections."),
+    "Jurisdiction": ("MEDIUM", "Disputes handled in their preferred court, not yours."),
+    "Arbitration": ("HIGH", "You waive your right to sue in court."),
+}
 
-NUM_LABELS = 8
-LABELS = [
-    ("Limitation of liability", "HIGH", "Company avoids responsibility for damages or losses."),
-    ("Unilateral termination", "HIGH", "They can close your account without reason."),
-    ("Unilateral change", "MEDIUM", "Terms can change without your consent."),
-    ("Content removal", "MEDIUM", "Your content can be deleted without notice."),
-    ("Contract by using", "LOW", "You agree just by visiting or using the site."),
-    ("Choice of law", "MEDIUM", "Foreign law applies instead of your local protections."),
-    ("Jurisdiction", "MEDIUM", "Disputes handled in their preferred court, not yours."),
-    ("Arbitration", "HIGH", "You waive your right to sue in court."),
-]
-
+# ─── Regex fallback ───
 PATTERNS = {
-    0: [r"not liable", r"shall not be (liable|responsible)", r"in no event.*liable", r"limitation of liability", r"without warranty", r"disclaim"],
-    1: [r"terminat.*at any time", r"suspend.*account.*without", r"we may (terminat|suspend|discontinu)", r"right to (terminat|suspend)"],
-    2: [r"sole discretion", r"reserves? the right to (modify|change|update|amend)", r"at any time.*without (prior )?notice", r"we may (modify|change|update)"],
-    3: [r"remove.*content.*without", r"right to remove", r"we may.*remove"],
-    4: [r"by (using|accessing).*you agree", r"continued use.*constitutes? acceptance"],
-    5: [r"governed by.*laws? of", r"shall be governed", r"laws of the state of"],
-    6: [r"exclusive jurisdiction", r"courts? of.*(california|delaware|new york|ireland|england)", r"submit to.*jurisdiction"],
-    7: [r"arbitrat", r"binding arbitration", r"waive.*right.*court", r"class action waiver"],
+    "Limitation of liability": [r"not liable", r"shall not be (liable|responsible)", r"in no event.*liable", r"limitation of liability", r"without warranty", r"disclaim"],
+    "Unilateral termination": [r"terminat.*at any time", r"suspend.*account.*without", r"we may (terminat|suspend|discontinu)", r"right to (terminat|suspend)"],
+    "Unilateral change": [r"sole discretion", r"reserves? the right to (modify|change|update|amend)", r"at any time.*without (prior )?notice", r"we may (modify|change|update)"],
+    "Content removal": [r"remove.*content.*without", r"right to remove", r"we may.*remove"],
+    "Contract by using": [r"by (using|accessing).*you agree", r"continued use.*constitutes? acceptance"],
+    "Choice of law": [r"governed by.*laws? of", r"shall be governed", r"laws of the state of"],
+    "Jurisdiction": [r"exclusive jurisdiction", r"courts? of.*(california|delaware|new york|ireland|england)", r"submit to.*jurisdiction"],
+    "Arbitration": [r"arbitrat", r"binding arbitration", r"waive.*right.*court", r"class action waiver"],
 }
 
+def classify_ml(text):
+    """Classify using the trained Legal-BERT model."""
+    if not ml_pipeline:
+        return classify_regex(text)
+    try:
+        preds = ml_pipeline(text, truncation=True, max_length=512)
+        results = []
+        for p in preds[0] if isinstance(preds[0], list) else preds:
+            if p["score"] > 0.5 and p["label"] in LABELS:
+                sev, desc = LABELS[p["label"]]
+                results.append({"name": p["label"], "severity": sev, "desc": desc, "confidence": round(p["score"], 2)})
+        return results
+    except Exception:
+        return classify_regex(text)
+
+def classify_regex(text):
+    """Fallback regex classifier."""
+    results = []
+    text_lower = text.lower()
+    for name, pats in PATTERNS.items():
+        for p in pats:
+            if re.search(p, text_lower):
+                sev, desc = LABELS[name]
+                results.append({"name": name, "severity": sev, "desc": desc, "confidence": 0.7})
+                break
+    return results
+
 def split_clauses(text):
     text = re.sub(r'\n{2,}', '\n', text.strip())
     parts = re.split(r'(?<=[.!?])\s+(?=[A-Z0-9(])|(?:\n)(?=\d+[.)]\s|\([a-z]\)\s)', text)
@@ -45,17 +86,11 @@ def analyze(text):
     sev_counts = {"HIGH": 0, "MEDIUM": 0, "LOW": 0}
 
     for clause in clauses:
-        clause_lower = clause.lower()
-        hits = []
-        for lid, pats in PATTERNS.items():
-            for p in pats:
-                if re.search(p, clause_lower):
-                    name, sev, desc = LABELS[lid]
-                    hits.append({"name": name, "severity": sev, "desc": desc})
-                    sev_counts[sev] += 1
-                    break
+        hits = classify_ml(clause)
         if hits:
             flagged.append({"text": clause, "hits": hits})
+            for h in hits:
+                sev_counts[h["severity"]] += 1
 
     total = len(clauses)
     risk = min(100, round((sev_counts["HIGH"] * 20 + sev_counts["MEDIUM"] * 10 + sev_counts["LOW"] * 5) / max(1, total) * 100))
@@ -66,7 +101,9 @@ def analyze(text):
     elif risk >= 10: grade = "B"
     else: grade = "A"
 
-    # Summary
+    engine = "Legal-BERT" if ml_pipeline else "Pattern matching"
+
+    # Build HTML
     summary = f"""<div style="font-family:system-ui,sans-serif;">
   <div style="border:1px solid #e4e4e7;border-radius:8px;padding:20px;margin-bottom:16px;">
     <div style="display:flex;justify-content:space-between;align-items:baseline;">
@@ -80,27 +117,28 @@ def analyze(text):
         'background:#f0fdf4;color:#15803d;'
       }">Grade {grade}</span>
     </div>
-    <p style="margin-top:8px;font-size:12px;color:#a1a1aa;">{total} clauses scanned · {len(flagged)} flagged · {sev_counts['HIGH']} high · {sev_counts['MEDIUM']} medium · {sev_counts['LOW']} low</p>
+    <p style="margin-top:8px;font-size:12px;color:#a1a1aa;">{total} clauses · {len(flagged)} flagged · {sev_counts['HIGH']} high · {sev_counts['MEDIUM']} medium · {sev_counts['LOW']} low · Engine: {engine}</p>
   </div>"""
 
     if not flagged:
         summary += '<div style="border:1px solid #e4e4e7;border-radius:8px;padding:24px;text-align:center;"><p style="font-size:14px;color:#71717a;">No unfair clauses found.</p></div>'
     else:
-        for i, item in enumerate(flagged):
+        for item in flagged:
             max_sev = max(item["hits"], key=lambda h: {"HIGH":3,"MEDIUM":2,"LOW":1}[h["severity"]])["severity"]
-            border_color = {"HIGH":"#fca5a5","MEDIUM":"#fcd34d","LOW":"#93c5fd"}[max_sev]
+            border = {"HIGH":"#fca5a5","MEDIUM":"#fcd34d","LOW":"#93c5fd"}[max_sev]
 
             tags = ""
             for h in item["hits"]:
-                tag_style = {"HIGH":"background:#fef2f2;color:#b91c1c;border:1px solid #fecaca;",
-                             "MEDIUM":"background:#fffbeb;color:#a16207;border:1px solid #fde68a;",
-                             "LOW":"background:#eff6ff;color:#1d4ed8;border:1px solid #bfdbfe;"}[h["severity"]]
-                tags += f'<span style="{tag_style}font-size:11px;font-weight:500;padding:1px 8px;border-radius:3px;margin-right:4px;">{h["name"]}</span>'
+                ts = {"HIGH":"background:#fef2f2;color:#b91c1c;border:1px solid #fecaca;",
+                      "MEDIUM":"background:#fffbeb;color:#a16207;border:1px solid #fde68a;",
+                      "LOW":"background:#eff6ff;color:#1d4ed8;border:1px solid #bfdbfe;"}[h["severity"]]
+                conf = f' ({h["confidence"]})' if h.get("confidence") and ml_pipeline else ""
+                tags += f'<span style="{ts}font-size:11px;font-weight:500;padding:1px 8px;border-radius:3px;margin-right:4px;">{h["name"]}{conf}</span>'
 
             descs = "".join(f'<p style="font-size:12px;color:#71717a;margin-top:4px;">{h["desc"]}</p>' for h in item["hits"])
-            preview = item["text"][:180] + ("..." if len(item["text"]) > 180 else "")
+            preview = item["text"][:200] + ("..." if len(item["text"]) > 200 else "")
 
-            summary += f'''<div style="border:1px solid #e4e4e7;border-left:3px solid {border_color};border-radius:8px;padding:14px;margin-bottom:8px;">
+            summary += f'''<div style="border:1px solid #e4e4e7;border-left:3px solid {border};border-radius:8px;padding:14px;margin-bottom:8px;">
     <p style="font-size:13px;color:#3f3f46;line-height:1.6;">{preview}</p>
     <div style="margin-top:8px;">{tags}</div>
     {descs}
@@ -109,6 +147,7 @@ def analyze(text):
     summary += "</div>"
     return summary, ""
 
+
 SPOTIFY = """By using the Spotify Service, you agree to be bound by these Terms of Use.
 
 Spotify may, in its sole discretion, modify or update these Terms of Service at any time without prior notice. Your continued use of the Service after any such changes constitutes your acceptance of the new Terms of Service.
@@ -153,7 +192,7 @@ with demo:
     scan_btn.click(fn=analyze, inputs=[text_input], outputs=[results_html, hidden])
     clear_btn.click(fn=lambda: ("", "", ""), outputs=[text_input, results_html, hidden])
 
-    gr.HTML('<p style="font-family:system-ui,sans-serif;font-size:11px;color:#a1a1aa;text-align:center;padding:16px 0;border-top:1px solid #f4f4f5;margin-top:16px;">Not legal advice. Based on CLAUDETTE taxonomy. <a href="https://huggingface.co/datasets/coastalcph/lex_glue" style="color:#71717a;">Dataset</a></p>')
+    gr.HTML('<p style="font-family:system-ui,sans-serif;font-size:11px;color:#a1a1aa;text-align:center;padding:16px 0;border-top:1px solid #f4f4f5;margin-top:16px;">Not legal advice. Model: Legal-BERT fine-tuned on CLAUDETTE. <a href="https://huggingface.co/gaurv007/clauseguard-legal-bert" style="color:#71717a;">Model</a> · <a href="https://huggingface.co/datasets/coastalcph/lex_glue" style="color:#71717a;">Dataset</a></p>')
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1 +1,4 @@
 gradio>=5.0
+transformers>=5.0
+torch
+numpy