Update app.py

app.py

@@ -1,407 +1,413 @@
-import streamlit as st
-import tempfile
-import os
-import re
-import io
-import json
-from typing import List, Dict, Tuple, Any, Optional
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from pypdf import PdfReader
-import docx
-import spacy
-import math
-
-# -------------------------
-# PAGE CONFIG
-# -------------------------
-st.set_page_config(page_title="ClauseWise – Granite 3.2 (2B) Legal Assistant", page_icon="⚖️", layout="wide")
-
-# -------------------------
-# MODEL SETUP
-# -------------------------
-MODEL_ID = "ibm-granite/granite-3.2-2b-instruct"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-@st.cache_resource
-def load_llm_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=DTYPE,
-        device_map="auto" if DEVICE == "cuda" else None
-    )
-    if DEVICE != "cuda":
-        model.to(DEVICE)
-    return tokenizer, model
-
-tokenizer, model = load_llm_model()
-
-nlp = spacy.load("en_core_web_sm")
-
-# -------------------------
-# HELPER FUNCTIONS
-# -------------------------
-def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
-    messages = []
-    if system_prompt:
-        messages.append({"role": "system", "content": system_prompt})
-    messages.append({"role": "user", "content": user_prompt})
-    try:
-        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    except Exception:
-        sys = f"[SYSTEM]\n{system_prompt}\n" if system_prompt else ""
-        usr = f"[USER]\n{user_prompt}\n[ASSISTANT]\n"
-        return sys + usr
-
-def llm_generate(system_prompt: str, user_prompt: str, max_new_tokens=512, temperature=0.3, top_p=0.9) -> str:
-    prompt = build_chat_prompt(system_prompt, user_prompt)
-    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-    with torch.inference_mode():
-        output_ids = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    if "[ASSISTANT]" in full_text:
-        return full_text.split("[ASSISTANT]")[-1].strip()
-    if full_text.startswith(prompt):
-        return full_text[len(prompt):].strip()
-    return full_text.strip()
-
-# -------------------------
-# DOCUMENT LOADING
-# -------------------------
-def load_text_from_pdf(file_obj) -> str:
-    reader = PdfReader(file_obj)
-    pages = []
-    for page in reader.pages:
-        try:
-            pages.append(page.extract_text() or "")
-        except Exception:
-            pages.append("")
-    return "\n".join(pages).strip()
-
-def load_text_from_docx(file_obj) -> str:
-    data = file_obj.read()
-    file_obj.seek(0)
-    f = io.BytesIO(data)
-    doc = docx.Document(f)
-    paras = [p.text for p in doc.paragraphs]
-    return "\n".join(paras).strip()
-
-def load_text_from_txt(file_obj) -> str:
-    data = file_obj.read()
-    if isinstance(data, bytes):
-        try:
-            data = data.decode("utf-8", errors="ignore")
-        except:
-            data = data.decode("latin-1", errors="ignore")
-    return str(data).strip()
-
-def load_document(file) -> str:
-    if not file:
-        return ""
-    name = (file.name or "").lower()
-    if name.endswith(".pdf"):
-        return load_text_from_pdf(file)
-    elif name.endswith(".docx"):
-        return load_text_from_docx(file)
-    elif name.endswith(".txt"):
-        return load_text_from_txt(file)
-    else:
-        try:
-            return load_text_from_pdf(file)
-        except:
-            pass
-        try:
-            return load_text_from_docx(file)
-        except:
-            pass
-        try:
-            return load_text_from_txt(file)
-        except:
-            pass
-        return ""
-
-def get_text_from_inputs(file, text):
-    file_text = load_document(file) if file else ""
-    final = (text or "").strip()
-    return file_text if len(file_text) > len(final) else final
-
-# -------------------------
-# CLAUSE PROCESSING
-# -------------------------
-CLAUSE_SPLIT_REGEX = re.compile(r"(?:(?:^\s*\d+(?:\.\d+)[.)]\s+)|(?:^\s[A-Z]\s*[.)]\s+)|(?:;?\s*\n))", re.MULTILINE)
-
-def split_into_clauses(text: str, min_len: int = 40) -> List[str]:
-    if not text:
-        return []
-    parts = re.split(CLAUSE_SPLIT_REGEX, text)
-    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]
-    seen = set()
-    unique = []
-    for c in clauses:
-        key = re.sub(r"\s+", " ", c.lower())
-        if key not in seen:
-            seen.add(key)
-            unique.append(c)
-    return unique
-
-def simplify_clause(clause: str) -> str:
-    system = "You are a legal assistant that rewrites clauses into plain English while preserving meaning."
-    user = f"Rewrite the following clause in plain English with bullet points for risks.\n\nClause:\n{clause}"
-    return llm_generate(system, user, max_new_tokens=400)
-
-def ner_entities(text: str) -> Dict[str, List[str]]:
-    if not text:
-        return {}
-    doc = nlp(text)
-    out: Dict[str, List[str]] = {}
-    for ent in doc.ents:
-        out.setdefault(ent.label_, []).append(ent.text)
-    out = {k: sorted(set(v)) for k, v in out.items()}
-    return out
-
-def extract_clauses(text: str) -> List[str]:
-    return split_into_clauses(text)
-
-# -------------------------
-# DOCUMENT CLASSIFICATION
-# -------------------------
-DOC_TYPES = [
-    "Non-Disclosure Agreement (NDA)",
-    "Lease Agreement",
-    "Employment Contract",
-    "Service Agreement",
-    "Sales Agreement",
-    "Consulting Agreement",
-    "End User License Agreement (EULA)",
-    "Terms of Service",
-]
-
-def classify_document(text: str) -> str:
-    system = "You are a legal document classifier. Choose the best matching document type."
-    labels = "\n".join(f"- {t}" for t in DOC_TYPES)
-    user = f"Classify the following document:\n{labels}\n\n{text[:5000]}"
-    resp = llm_generate(system, user, max_new_tokens=200)
-    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]
-    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
-
-# -------------------------
-# Negotiation Coach
-# -------------------------
-def negotiation_coach(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
-    system = "You are an AI negotiation coach."
-    user = (
-        "Propose 3 alternative versions ranked by acceptance rate in JSON.\n"
-        f"Clause:\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:
-        data = {"alternatives": []}
-        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."
-            })
-    return json.dumps(data, indent=2), data.get("alternatives", [])
-
-# -------------------------
-# Future Risk Predictor
-# -------------------------
-def future_risk_predictor(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
-    system = "Forecast future risks over 1–5 years."
-    user = f"Analyze clause and return JSON timeline.\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:
-        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"]
-            })
-    return json.dumps(data, indent=2), data["timeline"]
-
-# -------------------------
-# Fairness Balance Meter
-# -------------------------
-def fairness_balance_meter(clause: str) -> Tuple[str,int,str]:
-    system = "Evaluate clause fairness (0=Party A,50=balanced,100=Party B)."
-    user = f"Return JSON: score_0_100 and rationale.\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:
-        score,rationale=50,"Fallback balanced score."
-    return json.dumps({"score_0_100":score,"rationale":rationale,"notes":[]}, indent=2), score, rationale
-
-# -------------------------
-# Clause Battle Arena
-# -------------------------
-def clause_battle_arena(text_a: str, text_b: str) -> Tuple[str,str]:
-    system="Compare 2 contract drafts across categories."
-    user=f"Compare Document A vs Document B and return JSON.\nA:\n{text_a[:4000]}\nB:\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:
-        data={"rounds":[{"category":c,"winner":"Draw","rationale":"Fallback"} for c in ["Liability","Termination","IP","Payment","Confidentiality","Governing Law"]],
-              "overall_winner":"Draw","summary":"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
-
-# -------------------------
-# 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]]]:
-    system="Find hidden privacy traps and personal data."
-    user=f"Return JSON.\nText:\n{text[:6000]}"
-    resp=llm_generate(system,user,max_new_tokens=700)
-    try:data=json.loads(re.search(r"\{[\s\S]*\}",resp).group(0))
-    except:data={"data_categories":["Name","Email"],"sharing_parties":["Service Provider"],"processing_purposes":["Service delivery"],"risks":["Potential over-collection"],"recommendations":["Narrow purpose","Limit retention"]}
-    regex_hits={}
-    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
-
-# -------------------------
-# 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="Identify clauses likely to trigger disputes."
-    user=f"Return JSON of hotspots.\nClauses:\n{sample}"
-    resp=llm_generate(system,user,max_new_tokens=900)
-    try:data=json.loads(re.search(r"\{[\s\S]*\}",resp).group(0))
-    except: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
-
-# -------------------------
-# STREAMLIT UI
-# -------------------------
-st.title("ClauseWise – Granite 3.2 (2B) Legal Assistant")
-st.markdown("Upload a PDF/DOCX/TXT or paste text below. Tabs provide different legal analysis tools.")
-
-with st.sidebar:
-    uploaded_file = st.file_uploader("Upload PDF/DOCX/TXT (optional)", type=["pdf","docx","txt"])
-    pasted_text = st.text_area("Or paste text here", height=200)
-
-text_data = get_text_from_inputs(uploaded_file, pasted_text)
-
-tabs = st.tabs([
-    "Clause Simplification","Named Entity Recognition","Clause Extraction",
-    "Document Classification","Negotiation Coach","Future Risk Predictor",
-    "Fairness Balance Meter","Clause Battle Arena","Sensitive Data Sniffer","Litigation Risk Radar"
-])
-
-with tabs[0]:
-    clause_input = st.text_area("Clause (optional)", height=150)
-    if st.button("Simplify Clause", key="simplify"):
-        target = clause_input.strip() or text_data
-        st.text_area("Plain English Output", simplify_clause(target), height=250)
-
-with tabs[1]:
-    if st.button("Run NER", key="ner"):
-        st.json(ner_entities(text_data[:12000]))
-
-with tabs[2]:
-    if st.button("Extract Clauses", key="extract"):
-        clauses = extract_clauses(text_data)
-        st.dataframe([[c] for c in clauses], columns=["Clause"])
-
-with tabs[3]:
-    if st.button("Classify Document", key="classify"):
-        st.text_area("Predicted Type", classify_document(text_data))
-
-with tabs[4]:
-    negotiation_clause = st.text_area("Clause to Optimize", height=150)
-    if st.button("Suggest Alternatives", key="negotiation"):
-        pretty, alts = negotiation_coach(negotiation_clause.strip() or text_data)
-        st.json(json.loads(pretty))
-        table=[[a.get("rank",""),a.get("acceptance_rate_percent",""),a.get("title",""),a.get("clause_text",""),a.get("rationale","")] for a in alts]
-        st.dataframe(table, columns=["Rank","Acceptance %","Title","Clause Text","Rationale"])
-
-with tabs[5]:
-    risk_clause = st.text_area("Clause for Risk Prediction", height=150)
-    if st.button("Predict 1–5 Year Risks", key="risk"):
-        pretty, timeline = future_risk_predictor(risk_clause.strip() or text_data)
-        st.json(json.loads(pretty))
-        table=[[t.get("year",""),t.get("risk_score_0_100",""),"; ".join(t.get("key_risks",[])),"; ".join(t.get("mitigation",[]))] for t in timeline]
-        st.dataframe(table, columns=["Year","Risk Score (0–100)","Key Risks","Mitigation"])
-
-with tabs[6]:
-    fairness_clause = st.text_area("Clause", height=150)
-    if st.button("Compute Fairness", key="fairness"):
-        pretty, score, rationale = fairness_balance_meter(fairness_clause.strip() or text_data)
-        st.json(json.loads(pretty))
-        st.slider("Balance Score", min_value=0,max_value=100,value=score)
-        st.text_area("Rationale / Notes", rationale, height=100)
-
-with tabs[7]:
-    clause_a = st.text_area("Document A", height=150)
-    clause_b = st.text_area("Document B", height=150)
-    if st.button("Compare Clauses", key="battle"):
-        pretty, md = clause_battle_arena(clause_a.strip() or text_data, clause_b.strip() or text_data)
-        st.text_area("Battle JSON", pretty, height=300)
-        st.markdown(md)
-
-with tabs[8]:
-    if st.button("Scan for Sensitive Data", key="sensitive"):
-        pretty, hits = sensitive_data_sniffer(text_data)
-        st.text_area("Sensitive Data JSON", pretty, height=300)
-        st.json(hits)
-
-with tabs[9]:
-    if st.button("Identify Litigation Risk Hotspots", key="litigation"):
-        pretty, md = litigation_risk_radar(text_data)
-        st.text_area("Litigation JSON", pretty, height=300)
-        st.markdown(md)
+import streamlit as st
+import tempfile
+import os
+import re
+import io
+import json
+from typing import List, Dict, Tuple, Any, Optional
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from pypdf import PdfReader
+import docx
+import spacy
+import math
+
+# -------------------------
+# PAGE CONFIG
+# -------------------------
+st.set_page_config(page_title="ClauseWise – Granite 3.2 (2B) Legal Assistant", page_icon="⚖️", layout="wide")
+
+# -------------------------
+# MODEL SETUP
+# -------------------------
+MODEL_ID = "ibm-granite/granite-3.2-2b-instruct"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+@st.cache_resource
+def load_llm_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=DTYPE,
+        device_map="auto" if DEVICE == "cuda" else None
+    )
+    if DEVICE != "cuda":
+        model.to(DEVICE)
+    return tokenizer, model
+
+tokenizer, model = load_llm_model()
+
+from spacy.cli import download as spacy_download
+
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    spacy_download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
+
+# -------------------------
+# HELPER FUNCTIONS
+# -------------------------
+def build_chat_prompt(system_prompt: str, user_prompt: str) -> str:
+    messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+    messages.append({"role": "user", "content": user_prompt})
+    try:
+        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    except Exception:
+        sys = f"[SYSTEM]\n{system_prompt}\n" if system_prompt else ""
+        usr = f"[USER]\n{user_prompt}\n[ASSISTANT]\n"
+        return sys + usr
+
+def llm_generate(system_prompt: str, user_prompt: str, max_new_tokens=512, temperature=0.3, top_p=0.9) -> str:
+    prompt = build_chat_prompt(system_prompt, user_prompt)
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    if "[ASSISTANT]" in full_text:
+        return full_text.split("[ASSISTANT]")[-1].strip()
+    if full_text.startswith(prompt):
+        return full_text[len(prompt):].strip()
+    return full_text.strip()
+
+# -------------------------
+# DOCUMENT LOADING
+# -------------------------
+def load_text_from_pdf(file_obj) -> str:
+    reader = PdfReader(file_obj)
+    pages = []
+    for page in reader.pages:
+        try:
+            pages.append(page.extract_text() or "")
+        except Exception:
+            pages.append("")
+    return "\n".join(pages).strip()
+
+def load_text_from_docx(file_obj) -> str:
+    data = file_obj.read()
+    file_obj.seek(0)
+    f = io.BytesIO(data)
+    doc = docx.Document(f)
+    paras = [p.text for p in doc.paragraphs]
+    return "\n".join(paras).strip()
+
+def load_text_from_txt(file_obj) -> str:
+    data = file_obj.read()
+    if isinstance(data, bytes):
+        try:
+            data = data.decode("utf-8", errors="ignore")
+        except:
+            data = data.decode("latin-1", errors="ignore")
+    return str(data).strip()
+
+def load_document(file) -> str:
+    if not file:
+        return ""
+    name = (file.name or "").lower()
+    if name.endswith(".pdf"):
+        return load_text_from_pdf(file)
+    elif name.endswith(".docx"):
+        return load_text_from_docx(file)
+    elif name.endswith(".txt"):
+        return load_text_from_txt(file)
+    else:
+        try:
+            return load_text_from_pdf(file)
+        except:
+            pass
+        try:
+            return load_text_from_docx(file)
+        except:
+            pass
+        try:
+            return load_text_from_txt(file)
+        except:
+            pass
+        return ""
+
+def get_text_from_inputs(file, text):
+    file_text = load_document(file) if file else ""
+    final = (text or "").strip()
+    return file_text if len(file_text) > len(final) else final
+
+# -------------------------
+# CLAUSE PROCESSING
+# -------------------------
+CLAUSE_SPLIT_REGEX = re.compile(r"(?:(?:^\s*\d+(?:\.\d+)[.)]\s+)|(?:^\s[A-Z]\s*[.)]\s+)|(?:;?\s*\n))", re.MULTILINE)
+
+def split_into_clauses(text: str, min_len: int = 40) -> List[str]:
+    if not text:
+        return []
+    parts = re.split(CLAUSE_SPLIT_REGEX, text)
+    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]
+    seen = set()
+    unique = []
+    for c in clauses:
+        key = re.sub(r"\s+", " ", c.lower())
+        if key not in seen:
+            seen.add(key)
+            unique.append(c)
+    return unique
+
+def simplify_clause(clause: str) -> str:
+    system = "You are a legal assistant that rewrites clauses into plain English while preserving meaning."
+    user = f"Rewrite the following clause in plain English with bullet points for risks.\n\nClause:\n{clause}"
+    return llm_generate(system, user, max_new_tokens=400)
+
+def ner_entities(text: str) -> Dict[str, List[str]]:
+    if not text:
+        return {}
+    doc = nlp(text)
+    out: Dict[str, List[str]] = {}
+    for ent in doc.ents:
+        out.setdefault(ent.label_, []).append(ent.text)
+    out = {k: sorted(set(v)) for k, v in out.items()}
+    return out
+
+def extract_clauses(text: str) -> List[str]:
+    return split_into_clauses(text)
+
+# -------------------------
+# DOCUMENT CLASSIFICATION
+# -------------------------
+DOC_TYPES = [
+    "Non-Disclosure Agreement (NDA)",
+    "Lease Agreement",
+    "Employment Contract",
+    "Service Agreement",
+    "Sales Agreement",
+    "Consulting Agreement",
+    "End User License Agreement (EULA)",
+    "Terms of Service",
+]
+
+def classify_document(text: str) -> str:
+    system = "You are a legal document classifier. Choose the best matching document type."
+    labels = "\n".join(f"- {t}" for t in DOC_TYPES)
+    user = f"Classify the following document:\n{labels}\n\n{text[:5000]}"
+    resp = llm_generate(system, user, max_new_tokens=200)
+    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]
+    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
+
+# -------------------------
+# Negotiation Coach
+# -------------------------
+def negotiation_coach(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
+    system = "You are an AI negotiation coach."
+    user = (
+        "Propose 3 alternative versions ranked by acceptance rate in JSON.\n"
+        f"Clause:\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:
+        data = {"alternatives": []}
+        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."
+            })
+    return json.dumps(data, indent=2), data.get("alternatives", [])
+
+# -------------------------
+# Future Risk Predictor
+# -------------------------
+def future_risk_predictor(clause: str) -> Tuple[str, List[Dict[str, Any]]]:
+    system = "Forecast future risks over 1–5 years."
+    user = f"Analyze clause and return JSON timeline.\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:
+        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"]
+            })
+    return json.dumps(data, indent=2), data["timeline"]
+
+# -------------------------
+# Fairness Balance Meter
+# -------------------------
+def fairness_balance_meter(clause: str) -> Tuple[str,int,str]:
+    system = "Evaluate clause fairness (0=Party A,50=balanced,100=Party B)."
+    user = f"Return JSON: score_0_100 and rationale.\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:
+        score,rationale=50,"Fallback balanced score."
+    return json.dumps({"score_0_100":score,"rationale":rationale,"notes":[]}, indent=2), score, rationale
+
+# -------------------------
+# Clause Battle Arena
+# -------------------------
+def clause_battle_arena(text_a: str, text_b: str) -> Tuple[str,str]:
+    system="Compare 2 contract drafts across categories."
+    user=f"Compare Document A vs Document B and return JSON.\nA:\n{text_a[:4000]}\nB:\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:
+        data={"rounds":[{"category":c,"winner":"Draw","rationale":"Fallback"} for c in ["Liability","Termination","IP","Payment","Confidentiality","Governing Law"]],
+              "overall_winner":"Draw","summary":"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
+
+# -------------------------
+# 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]]]:
+    system="Find hidden privacy traps and personal data."
+    user=f"Return JSON.\nText:\n{text[:6000]}"
+    resp=llm_generate(system,user,max_new_tokens=700)
+    try:data=json.loads(re.search(r"\{[\s\S]*\}",resp).group(0))
+    except:data={"data_categories":["Name","Email"],"sharing_parties":["Service Provider"],"processing_purposes":["Service delivery"],"risks":["Potential over-collection"],"recommendations":["Narrow purpose","Limit retention"]}
+    regex_hits={}
+    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
+
+# -------------------------
+# 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="Identify clauses likely to trigger disputes."
+    user=f"Return JSON of hotspots.\nClauses:\n{sample}"
+    resp=llm_generate(system,user,max_new_tokens=900)
+    try:data=json.loads(re.search(r"\{[\s\S]*\}",resp).group(0))
+    except: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
+
+# -------------------------
+# STREAMLIT UI
+# -------------------------
+st.title("ClauseWise – Granite 3.2 (2B) Legal Assistant")
+st.markdown("Upload a PDF/DOCX/TXT or paste text below. Tabs provide different legal analysis tools.")
+
+with st.sidebar:
+    uploaded_file = st.file_uploader("Upload PDF/DOCX/TXT (optional)", type=["pdf","docx","txt"])
+    pasted_text = st.text_area("Or paste text here", height=200)
+
+text_data = get_text_from_inputs(uploaded_file, pasted_text)
+
+tabs = st.tabs([
+    "Clause Simplification","Named Entity Recognition","Clause Extraction",
+    "Document Classification","Negotiation Coach","Future Risk Predictor",
+    "Fairness Balance Meter","Clause Battle Arena","Sensitive Data Sniffer","Litigation Risk Radar"
+])
+
+with tabs[0]:
+    clause_input = st.text_area("Clause (optional)", height=150)
+    if st.button("Simplify Clause", key="simplify"):
+        target = clause_input.strip() or text_data
+        st.text_area("Plain English Output", simplify_clause(target), height=250)
+
+with tabs[1]:
+    if st.button("Run NER", key="ner"):
+        st.json(ner_entities(text_data[:12000]))
+
+with tabs[2]:
+    if st.button("Extract Clauses", key="extract"):
+        clauses = extract_clauses(text_data)
+        st.dataframe([[c] for c in clauses], columns=["Clause"])
+
+with tabs[3]:
+    if st.button("Classify Document", key="classify"):
+        st.text_area("Predicted Type", classify_document(text_data))
+
+with tabs[4]:
+    negotiation_clause = st.text_area("Clause to Optimize", height=150)
+    if st.button("Suggest Alternatives", key="negotiation"):
+        pretty, alts = negotiation_coach(negotiation_clause.strip() or text_data)
+        st.json(json.loads(pretty))
+        table=[[a.get("rank",""),a.get("acceptance_rate_percent",""),a.get("title",""),a.get("clause_text",""),a.get("rationale","")] for a in alts]
+        st.dataframe(table, columns=["Rank","Acceptance %","Title","Clause Text","Rationale"])
+
+with tabs[5]:
+    risk_clause = st.text_area("Clause for Risk Prediction", height=150)
+    if st.button("Predict 1–5 Year Risks", key="risk"):
+        pretty, timeline = future_risk_predictor(risk_clause.strip() or text_data)
+        st.json(json.loads(pretty))
+        table=[[t.get("year",""),t.get("risk_score_0_100",""),"; ".join(t.get("key_risks",[])),"; ".join(t.get("mitigation",[]))] for t in timeline]
+        st.dataframe(table, columns=["Year","Risk Score (0–100)","Key Risks","Mitigation"])
+
+with tabs[6]:
+    fairness_clause = st.text_area("Clause", height=150)
+    if st.button("Compute Fairness", key="fairness"):
+        pretty, score, rationale = fairness_balance_meter(fairness_clause.strip() or text_data)
+        st.json(json.loads(pretty))
+        st.slider("Balance Score", min_value=0,max_value=100,value=score)
+        st.text_area("Rationale / Notes", rationale, height=100)
+
+with tabs[7]:
+    clause_a = st.text_area("Document A", height=150)
+    clause_b = st.text_area("Document B", height=150)
+    if st.button("Compare Clauses", key="battle"):
+        pretty, md = clause_battle_arena(clause_a.strip() or text_data, clause_b.strip() or text_data)
+        st.text_area("Battle JSON", pretty, height=300)
+        st.markdown(md)
+
+with tabs[8]:
+    if st.button("Scan for Sensitive Data", key="sensitive"):
+        pretty, hits = sensitive_data_sniffer(text_data)
+        st.text_area("Sensitive Data JSON", pretty, height=300)
+        st.json(hits)
+
+with tabs[9]:
+    if st.button("Identify Litigation Risk Hotspots", key="litigation"):
+        pretty, md = litigation_risk_radar(text_data)
+        st.text_area("Litigation JSON", pretty, height=300)
+        st.markdown(md)