Update app.py

app.py

@@ -3,18 +3,17 @@ import json
 import math
 import re
 import io
-import asyncio
 from typing import List, Dict, Tuple, Optional, Any
 
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from pypdf import PdfReader
 import docx
 import spacy
 import gradio as gr
 
 # -----------------------------
-# Model: IBM Granite 3.2 2B Instruct
+# Model: IBM Granite 3.2 2B Instruct (you can swap this for a local HF model)
 # -----------------------------
 MODEL_ID = "ibm-granite/granite-3.2-2b-instruct"
 
@@ -36,7 +35,7 @@ if DEVICE != "cuda":
 nlp = spacy.load("en_core_web_sm")
 
 # -----------------------------
-# Helper: chat templating for Granite or fallback
+# Chat templating helper
 # -----------------------------
 def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
     messages = []
@@ -46,7 +45,6 @@ def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
     try:
         return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     except Exception:
-        # Fallback: simple concatenation
         sys = f"[SYSTEM]\n{system_prompt}\n" if system_prompt else ""
         usr = f"[USER]\n{user_prompt}\n[ASSISTANT]\n"
         return sys + usr
@@ -54,13 +52,7 @@ def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
 # -----------------------------
 # LLM generation
 # -----------------------------
-def llm_generate(
-    system_prompt: str,
-    user_prompt: str,
-    max_new_tokens: int = 512,
-    temperature: float = 0.3,
-    top_p: float = 0.9
-) -> str:
+def llm_generate(system_prompt: str, user_prompt: str, max_new_tokens: int = 512, temperature: float = 0.3, top_p: float = 0.9) -> str:
     prompt = build_chat_prompt(system_prompt, user_prompt)
     inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     with torch.inference_mode():
@@ -73,10 +65,8 @@ def llm_generate(
             pad_token_id=tokenizer.eos_token_id
         )
     full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    # Try to extract only assistant part if chat template used
     if "[ASSISTANT]" in full_text:
         return full_text.split("[ASSISTANT]")[-1].strip()
-    # Otherwise, remove the prompt prefix
     if full_text.startswith(prompt):
         return full_text[len(prompt):].strip()
     return full_text.strip()
@@ -95,7 +85,6 @@ def load_text_from_pdf(file_obj) -> str:
     return "\n".join(pages).strip()
 
 def load_text_from_docx(file_obj) -> str:
-    # file_obj is a temporary file-like object; need to read into BytesIO for python-docx
     data = file_obj.read()
     file_obj.seek(0)
     f = io.BytesIO(data)
@@ -123,7 +112,6 @@ def load_document(file: Optional[gr.File]) -> str:
     elif name.endswith(".txt"):
         return load_text_from_txt(file)
     else:
-        # Try all in order
         try:
             return load_text_from_pdf(file)
         except Exception:
@@ -149,13 +137,10 @@ CLAUSE_SPLIT_REGEX = re.compile(
 def split_into_clauses(text: str, min_len: int = 40) -> List[str]:
     if not text:
         return []
-    # First try structured numbering/bullets
     parts = re.split(CLAUSE_SPLIT_REGEX, text)
-    # Fallback: sentence-like splits if too few
     if len(parts) < 2:
         parts = re.split(r"(?<=[.;])\s+\n?\s*", text)
     clauses = [p.strip() for p in parts if len(p.strip()) >= min_len]
-    # Deduplicate near-identical snippets
     seen = set()
     unique = []
     for c in clauses:
@@ -166,15 +151,15 @@ def split_into_clauses(text: str, min_len: int = 40) -> List[str]:
     return unique
 
 # -----------------------------
-# Feature: Clause Simplification / Plain English
+# Simplify clause
 # -----------------------------
 def simplify_clause(clause: str) -> str:
-    system = "You are a legal assistant that rewrites clauses into plain, layman-friendly English while preserving legal meaning."
-    user = f"Rewrite the following clause in plain English, preserving intent. Highlight any risks with bullet points at the end.\n\nClause:\n{clause}"
+    system = "You are a legal assistant that rewrites clauses into plain English while preserving legal meaning."
+    user = f"Rewrite this clause in plain English:\n\n{clause}"
     return llm_generate(system, user, max_new_tokens=400)
 
 # -----------------------------
-# Feature: Named Entity Recognition (NER)
+# Named Entity Recognition
 # -----------------------------
 def ner_entities(text: str) -> Dict[str, List[str]]:
     if not text:
@@ -183,18 +168,16 @@ def ner_entities(text: str) -> Dict[str, List[str]]:
     out: Dict[str, List[str]] = {}
     for ent in doc.ents:
         out.setdefault(ent.label_, []).append(ent.text)
-    # Deduplicate
-    out = {k: sorted(set(v)) for k, v in out.items()}
-    return out
+    return {k: sorted(set(v)) for k, v in out.items()}
 
 # -----------------------------
-# Feature: Clause Extraction and Breakdown
+# Clause extraction
 # -----------------------------
 def extract_clauses(text: str) -> List[str]:
     return split_into_clauses(text)
 
 # -----------------------------
-# Feature: Document Type Classification (LLM zero-shot)
+# Document classification
 # -----------------------------
 DOC_TYPES = [
     "Non-Disclosure Agreement (NDA)",
@@ -208,215 +191,21 @@ DOC_TYPES = [
 ]
 
 def classify_document(text: str) -> str:
-    system = "You are a legal document classifier. Choose the single best-matching document type from the provided list."
+    system = "You are a legal document classifier."
     labels = "\n".join(f"- {t}" for t in DOC_TYPES)
-    user = f"Classify the following document into one of these types:\n{labels}\n\nDocument:\n{text[:5000]}"
+    user = f"Classify the document into one of these types:\n{labels}\n\nDocument:\n{text[:5000]}"
     resp = llm_generate(system, user, max_new_tokens=200)
-    # Try to pick the closest label
-    scores = {t: (1.0 if t.lower() in resp.lower() else 0.0) for t in DOC_TYPES}
-    best = max(scores.items(), key=lambda kv: kv[1])[0]
-    # Fallback: heuristic keyword match
-    if scores[best] == 0.0:
-        lower = text.lower()
-        if "confidential" in lower or "non-disclosure" in lower or "nda" in lower:
-            best = "Non-Disclosure Agreement (NDA)"
-        elif "lease" in lower or "tenant" in lower or "landlord" in lower:
-            best = "Lease Agreement"
-        elif "employment" in lower or "employee" in lower or "employer" in lower:
-            best = "Employment Contract"
-        elif "services" in lower or "service" in lower or "statement of work" in lower:
-            best = "Service Agreement"
-    return best
+    for t in DOC_TYPES:
+        if t.lower() in resp.lower():
+            return t
+    return "Unclassified"
 
 # -----------------------------
-# Feature: Negotiation Coach (3 alternatives with acceptance rates)
-# -----------------------------
-def negotiation_coach(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
-    system = "You are an AI negotiation coach for contracts."
-    user = (
-        "Given the clause below, propose 3 alternative versions ranked by expected acceptance rate. "
-        "Provide JSON with fields: alternatives: [ {rank, acceptance_rate_percent, title, clause_text, rationale} ]. "
-        "Rank 1 is highest acceptance rate. Keep acceptance_rate_percent as integer. "
-        f"\n\nClause:\n{clause}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=700)
-    # Try to extract JSON
-    data = None
-    try:
-        json_str = re.search(r"\{[\s\S]*\}", resp).group(0)
-        data = json.loads(json_str)
-    except Exception:
-        # Try to reconstruct minimal structure
-        data = {"alternatives": []}
-        # heuristic extraction
-        alts = re.split(r"\n\s*\d+[.)]\s*", resp)
-        for i, chunk in enumerate(alts[1:4], start=1):
-            data["alternatives"].append({
-                "rank": i,
-                "acceptance_rate_percent": max(50, 90 - (i-1)*10),
-                "title": f"Alternative {i}",
-                "clause_text": chunk.strip()[:800],
-                "rationale": "Heuristic parse from model response."
-            })
-    pretty = json.dumps(data, indent=2)
-    return pretty, data.get("alternatives", [])
-
-# -----------------------------
-# Feature: Future Risk Predictor (1–5+ years timeline)
-# -----------------------------
-def future_risk_predictor(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
-    system = "You analyze contractual clauses and forecast future risks over time."
-    user = (
-        "Analyze the clause and forecast risks over the next 1 to 5 years. "
-        "Return strict JSON: {timeline: [ {year: int, risk_score_0_100: int, key_risks: [str], mitigation: [str]} ]}. "
-        "risk_score_0_100 is an integer. Keep the list length between 5 and 6."
-        f"\n\nClause:\n{clause}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=700)
-    data = None
-    try:
-        json_str = re.search(r"\{[\s\S]*\}", resp).group(0)
-        data = json.loads(json_str)
-    except Exception:
-        data = {"timeline": []}
-        for y in range(1, 6):
-            data["timeline"].append({
-                "year": y,
-                "risk_score_0_100": min(95, 40 + y*8),
-                "key_risks": ["Heuristic timeline due to JSON parse fallback."],
-                "mitigation": ["Seek legal review", "Adjust clause terms", "Add notice/cure period"]
-            })
-    pretty = json.dumps(data, indent=2)
-    return pretty, data["timeline"]
-
-# -----------------------------
-# Feature: Fairness Balance Meter (power distribution)
-# -----------------------------
-def fairness_balance_meter(clause: str) -> Tuple[str, int, str]:
-    system = "You evaluate which party a clause favors on a 0-100 scale (0=Party A heavily favored, 50=balanced, 100=Party B heavily favored)."
-    user = (
-        "Return strict JSON: {score_0_100: int, rationale: str, notes: [str]}. "
-        "Do not include anything else."
-        f"\n\nClause:\n{clause}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=400)
-    try:
-        data = json.loads(re.search(r"\{[\s\S]*\}", resp).group(0))
-        score = int(data.get("score_0_100", 50))
-        rationale = data.get("rationale", "")
-    except Exception:
-        score, rationale = 50, "Fallback balanced score due to JSON parse."
-        data = {"score_0_100": score, "rationale": rationale, "notes": []}
-    pretty = json.dumps(data, indent=2)
-    return pretty, score, rationale
-
-# -----------------------------
-# Feature: Clause Battle Arena (head-to-head)
-# -----------------------------
-def clause_battle_arena(text_a: str, text_b: str) -> Tuple[str, str]:
-    system = "You compare two contract drafts across objective criteria and declare an overall winner."
-    user = (
-        "Compare Document A vs Document B across: Liability, Termination, IP, Payment, Confidentiality, Governing Law. "
-        "Return JSON: {rounds: [ {category, winner: 'A'|'B'|'Draw', rationale} ], overall_winner: 'A'|'B'|'Draw', summary: str}.\n"
-        f"Document A:\n{text_a[:4000]}\n\nDocument B:\n{text_b[:4000]}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=900)
-    try:
-        data = json.loads(re.search(r"\{[\s\S]*\}", resp).group(0))
-    except Exception:
-        data = {
-            "rounds": [
-                {"category": "Liability", "winner": "Draw", "rationale": "Fallback"},
-                {"category": "Termination", "winner": "Draw", "rationale": "Fallback"},
-                {"category": "IP", "winner": "Draw", "rationale": "Fallback"},
-                {"category": "Payment", "winner": "Draw", "rationale": "Fallback"},
-                {"category": "Confidentiality", "winner": "Draw", "rationale": "Fallback"},
-                {"category": "Governing Law", "winner": "Draw", "rationale": "Fallback"},
-            ],
-            "overall_winner": "Draw",
-            "summary": "JSON parse fallback."
-        }
-    pretty = json.dumps(data, indent=2)
-    rounds_md = "\n".join([f"- {r['category']}: {r['winner']} — {r.get('rationale','')}" for r in data.get("rounds", [])])
-    md = f"Overall Winner: {data.get('overall_winner','Draw')}\n\nRounds:\n{rounds_md}\n\nSummary:\n{data.get('summary','')}"
-    return pretty, md
-
-# -----------------------------
-# Feature: Sensitive Data Sniffer
-# -----------------------------
-PII_REGEXES = {
-    "Email": r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}",
-    "Phone": r"\+?\d[\d\-\s]{7,}\d",
-    "SSN (US)": r"\b\d{3}-\d{2}-\d{4}\b",
-    "Credit Card": r"\b(?:\d[ -]*?){13,16}\b",
-}
-
-def sensitive_data_sniffer(text: str) -> Tuple[str, Dict[str, List[str]]]:
-    # LLM-based identification of privacy traps plus regex PII
-    system = "You find hidden privacy traps in legal text and list personal data categories being shared or processed."
-    user = (
-        "Return strict JSON: {data_categories: [str], sharing_parties: [str], processing_purposes: [str], risks: [str], recommendations: [str]}.\n"
-        f"Text:\n{text[:6000]}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=700)
-    data = None
-    try:
-        data = json.loads(re.search(r"\{[\s\S]*\}", resp).group(0))
-    except Exception:
-        data = {
-            "data_categories": ["Name", "Email"],
-            "sharing_parties": ["Service Provider"],
-            "processing_purposes": ["Service delivery"],
-            "risks": ["Potential over-collection"],
-            "recommendations": ["Narrow purpose", "Limit retention"]
-        }
-    # Regex-based PII findings
-    regex_hits: Dict[str, List[str]] = {}
-    for label, pattern in PII_REGEXES.items():
-        hits = re.findall(pattern, text or "", flags=re.IGNORECASE)
-        if hits:
-            regex_hits[label] = sorted(set([h.strip() for h in hits]))
-    pretty = json.dumps({"llm": data, "regex_hits": regex_hits}, indent=2)
-    return pretty, regex_hits
-
-# -----------------------------
-# Feature: Litigation Risk Radar
-# -----------------------------
-def litigation_risk_radar(text: str) -> Tuple[str, str]:
-    clauses = split_into_clauses(text)
-    sample = "\n\n".join(clauses[:8]) if clauses else text[:4000]
-    system = "You identify clauses most likely to trigger disputes or litigation and provide sample dispute scenarios."
-    user = (
-        "Analyze the clauses and return JSON: {hotspots: [ {clause_excerpt, risk_level: 'Low'|'Medium'|'High', why, sample_dispute_scenario} ]}.\n"
-        f"Clauses:\n{sample}"
-    )
-    resp = llm_generate(system, user, max_new_tokens=900)
-    try:
-        data = json.loads(re.search(r"\{[\s\S]*\}", resp).group(0))
-    except Exception:
-        data = {
-            "hotspots": [
-                {
-                    "clause_excerpt": (clauses[0][:280] if clauses else text[:280]),
-                    "risk_level": "Medium",
-                    "why": "Ambiguous obligations.",
-                    "sample_dispute_scenario": "Party A alleges non-performance due to unclear milestones."
-                }
-            ]
-        }
-    pretty = json.dumps(data, indent=2)
-    md = "\n".join([
-        f"- [{h.get('risk_level','Medium')}] {h.get('clause_excerpt','')}\n  Why: {h.get('why','')}\n  Scenario: {h.get('sample_dispute_scenario','')}"
-        for h in data.get("hotspots", [])
-    ])
-    return pretty, md
-
-# -----------------------------
-# Glue: Input handling (upload or paste)
+# Input Handling
 # -----------------------------
 def get_text_from_inputs(file: Optional[gr.File], text: str) -> str:
     file_text = load_document(file) if file else ""
     final = (text or "").strip()
     if len(file_text) > len(final):
         return file_text
-    return final
+    return final