Update app.py

app.py

@@ -1,178 +1,225 @@
 import streamlit as st
-
-st.set_page_config(
-    page_title="ClauseWise — NDA Legal Assistant",
-    page_icon="📄",
-    layout="wide"
-)
-import os
-import json
-from pypdf import PdfReader
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from multilingual import UI_TEXT, translate_text
+import torch
+
+# Import utility functions
 from util import (
     extract_text_from_pdf,
+    extract_text_from_txt,
+    extract_text_from_docx,
+    detect_if_nda,
     split_into_clauses,
     simplify_clause,
     extract_entities,
-    generate_risks,
-    fairness_score,
-    alternative_clauses
+    extract_top_risks,
+    compute_fairness,
+    get_alternative_clauses
 )
 
-# ------------------------------------------------------------
-# ✅ Load Local Model (DistilGPT2)
-# ------------------------------------------------------------
+# Multilingual UI
+from multilingual import UI_TEXT, translate_text
+
+
+# -----------------------------------
+# ✅ INITIAL STREAMLIT PAGE CONFIG
+# -----------------------------------
+st.set_page_config(page_title="ClauseWise", layout="wide")
+
+
+# -----------------------------------
+# ✅ LOAD LOCAL SMALL MODEL (HF SAFE)
+# -----------------------------------
 @st.cache_resource
-def load_model():
+def load_llm():
     tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
     model = AutoModelForCausalLM.from_pretrained("distilgpt2")
     return tokenizer, model
 
-tokenizer, model = load_model()
 
-# ------------------------------------------------------------
-# ✅ Chatbot Response
-# ------------------------------------------------------------
-def chat_response(prompt):
+tokenizer, model = load_llm()
+
+
+# -----------------------------------
+# ✅ LLM GENERATION FUNCTION
+# -----------------------------------
+def chat_ai(prompt):
     inputs = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=180, num_return_sequences=1)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return text
+    outputs = model.generate(
+        inputs,
+        max_new_tokens=150,
+        temperature=0.7,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# ------------------------------------------------------------
-# ✅ Streamlit Layout
-# ------------------------------------------------------------
-st.set_page_config(page_title="ClauseWise", layout="wide")
 
-# ------------------------------------------------------------
-# ✅ Session State
-# ------------------------------------------------------------
+# -----------------------------------
+# ✅ SESSION STATE
+# -----------------------------------
+if "language" not in st.session_state:
+    st.session_state.language = "en"
+
+if "uploaded_text" not in st.session_state:
+    st.session_state.uploaded_text = ""
+
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
-# ------------------------------------------------------------
-# ✅ Language Selector
-# ------------------------------------------------------------
-lang = st.selectbox(
-    "🌐 Select Language",
+
+# -----------------------------------
+# ✅ LANGUAGE SELECTION
+# -----------------------------------
+st.sidebar.title("🌐 Language")
+lang = st.sidebar.selectbox(
+    "Choose your language",
     ["en", "hi", "ta", "te", "kn"],
-    format_func=lambda x: {
-        "en": "English",
-        "hi": "Hindi",
-        "ta": "Tamil",
-        "te": "Telugu",
-        "kn": "Kannada"
-    }[x]
+    index=["en", "hi", "ta", "te", "kn"].index(st.session_state.language)
 )
 
-st.markdown("### **📄 ClauseWise — Legal Document Assistant**")
-st.info("Educational information only. Not legal advice.")
+st.session_state.language = lang
+T = UI_TEXT[lang]
 
-# ------------------------------------------------------------
-# ✅ File Upload
-# ------------------------------------------------------------
-st.markdown(f"### {UI_TEXT['upload_title'][lang]}")
+
+# -----------------------------------
+# ✅ PAGE TITLE
+# -----------------------------------
+st.title("📝 ClauseWise — Multilingual NDA Legal Assistant")
+
+
+# -----------------------------------
+# ✅ FILE UPLOAD SECTION
+# -----------------------------------
+st.header(T["upload_section"])
 
 uploaded_file = st.file_uploader(
-    UI_TEXT["upload_instruction"][lang],
+    T["upload_instruction"],
     type=["pdf", "txt", "docx"]
 )
 
-text_data = ""
+if uploaded_file:
+    file_ext = uploaded_file.name.split(".")[-1].lower()
+
+    if file_ext == "pdf":
+        text = extract_text_from_pdf(uploaded_file)
+    elif file_ext == "txt":
+        text = extract_text_from_txt(uploaded_file)
+    elif file_ext == "docx":
+        text = extract_text_from_docx(uploaded_file)
+    else:
+        st.error(T["unsupported"])
+        st.stop()
 
-if uploaded_file is not None:
-    # ✅ Extract text depending on file type
-    if uploaded_file.type == "application/pdf":
-        text_data = extract_text_from_pdf(uploaded_file)
-    elif uploaded_file.type == "text/plain":
-        text_data = uploaded_file.read().decode("utf-8")
-    elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
-        import docx
-        doc = docx.Document(uploaded_file)
-        text_data = "\n".join([p.text for p in doc.paragraphs])
+    st.session_state.uploaded_text = text
 
     # ✅ NDA Detection
-    if "confidential" not in text_data.lower():
-        st.error(UI_TEXT["error_not_nda"][lang])
+    if detect_if_nda(text):
+        st.success(T["nda_detected"])
     else:
-        st.success(UI_TEXT["success_nda"][lang])
-
-        # ------------------------------------------------------------
-        # ✅ Document Tabs
-        # ------------------------------------------------------------
-        tab1, tab2, tab3, tab4, tab5, tab6 = st.tabs([
-            UI_TEXT["tab_clauses"][lang],
-            UI_TEXT["tab_risks"][lang],
-            UI_TEXT["tab_fairness"][lang],
-            UI_TEXT["tab_entities"][lang],
-            UI_TEXT["tab_alternatives"][lang],
-            UI_TEXT["tab_chat"][lang]
-        ])
-
-        clauses = split_into_clauses(text_data)
-
-        # ------------------------------------------------------------
-        # ✅ Clause Simplification
-        # ------------------------------------------------------------
-        with tab1:
-            st.subheader(UI_TEXT["clause_simplify"][lang])
-            mode = st.selectbox(
-                UI_TEXT["choose_mode"][lang],
-                ["ELI5", "Simple", "Professional"]
-            )
-            for cl in clauses:
-                st.markdown(f"### Clause:\n{cl}")
-                st.markdown("**Explanation:**")
-                st.write(simplify_clause(cl, mode))
-
-        # ------------------------------------------------------------
-        # ✅ Risk Analysis
-        # ------------------------------------------------------------
-        with tab2:
-            st.subheader(UI_TEXT["risk_title"][lang])
-            risks = generate_risks(text_data)
-            for r in risks:
-                st.warning(f"⚠️ {r}")
-
-        # ------------------------------------------------------------
-        # ✅ Fairness Meter
-        # ------------------------------------------------------------
-        with tab3:
-            st.subheader(UI_TEXT["fairness_title"][lang])
-            emp, comp = fairness_score(text_data)
-            st.write(f"**{UI_TEXT['your_position'][lang]}:** {emp}")
-            st.write(f"**{UI_TEXT['company_position'][lang]}:** {comp}")
-
-        # ------------------------------------------------------------
-        # ✅ Entity Extraction
-        # ------------------------------------------------------------
-        with tab4:
-            st.subheader(UI_TEXT["entities_title"][lang])
-            ents = extract_entities(text_data)
-            st.json(ents)
-
-        # ------------------------------------------------------------
-        # ✅ Alternative Clauses
-        # ------------------------------------------------------------
-        with tab5:
-            st.subheader(UI_TEXT["alt_title"][lang])
-            for cl in clauses:
-                st.write(f"📄 **Original:** {cl}")
-                st.write(f"✅ **Alternative:** {alternative_clauses(cl)}")
-
-        # ------------------------------------------------------------
-        # ✅ Legal Chat Assistant
-        # ------------------------------------------------------------
-        with tab6:
-            st.subheader(UI_TEXT["chat_title"][lang])
-            user_input = st.text_input(UI_TEXT["chat_placeholder"][lang])
-
-            if st.button("Send"):
-                if user_input.strip():
-                    reply = chat_response(user_input)
-                    st.session_state.chat_history.append(("You", user_input))
-                    st.session_state.chat_history.append(("ClauseWise", reply))
-
-            for sender, msg in st.session_state.chat_history:
-                st.write(f"**{sender}:** {msg}")
+        st.warning(T["not_nda"])
+
+
+# =====================================================================
+# ✅ TABS
+# =====================================================================
+if st.session_state.uploaded_text.strip():
+    tabs = st.tabs([
+        T["tab_simplify"],
+        T["tab_risks"],
+        T["tab_fairness"],
+        T["tab_entities"],
+        T["tab_alternatives"],
+        T["tab_chat"]
+    ])
+
+    text = st.session_state.uploaded_text
+
+    # -----------------------------------------------------------
+    # ✅ TAB 1 — Clause Simplification
+    # -----------------------------------------------------------
+    with tabs[0]:
+        st.subheader(T["simplify_title"])
+
+        mode = st.selectbox(
+            T["choose_simplify_mode"],
+            ["simple", "eli5", "pro"],
+            format_func=lambda m: {
+                "simple": T["simple"],
+                "eli5": T["eli5"],
+                "pro": T["pro"]
+            }[m]
+        )
+
+        clauses = split_into_clauses(text)
+
+        for c in clauses:
+            st.write("### 📌 " + T["clause"])
+            st.write(c)
+
+            simp = simplify_clause(c, mode=mode)
+            st.info(simp)
+            st.divider()
+
+    # -----------------------------------------------------------
+    # ✅ TAB 2 — Risk Analysis
+    # -----------------------------------------------------------
+    with tabs[1]:
+        st.subheader(T["risk_title"])
+        risks = extract_top_risks(text)
+
+        for r in risks:
+            st.warning("⚠️ " + r)
+
+    # -----------------------------------------------------------
+    # ✅ TAB 3 — Fairness Meter
+    # -----------------------------------------------------------
+    with tabs[2]:
+        st.subheader(T["fairness_title"])
+
+        fairness = compute_fairness(text)
+
+        st.metric(label=T["your_position"], value=f"{fairness['your_position']} / 100")
+        st.metric(label=T["their_position"], value=f"{fairness['their_position']} / 100")
+
+    # -----------------------------------------------------------
+    # ✅ TAB 4 — Entities
+    # -----------------------------------------------------------
+    with tabs[3]:
+        st.subheader(T["entities_title"])
+
+        ents = extract_entities(text)
+        st.write("### " + T["names"])
+        st.write(ents["names"])
+
+        st.write("### " + T["dates"])
+        st.write(ents["dates"])
+
+    # -----------------------------------------------------------
+    # ✅ TAB 5 — Alternative Clauses
+    # -----------------------------------------------------------
+    with tabs[4]:
+        st.subheader(T["alt_title"])
+
+        alts = get_alternative_clauses(text)
+        for a in alts:
+            st.info("✅ " + a)
+
+    # -----------------------------------------------------------
+    # ✅ TAB 6 — Legal Chat Assistant
+    # -----------------------------------------------------------
+    with tabs[5]:
+        st.subheader(T["chat_title"])
+
+        # show chat history
+        for item in st.session_state.chat_history:
+            st.chat_message(item["role"]).write(item["content"])
+
+        # chat input
+        user_input = st.chat_input(T["chat_placeholder"])
+        if user_input:
+            st.session_state.chat_history.append({"role": "user", "content": user_input})
+            st.chat_message("user").write(user_input)
+
+            # model response
+            response = chat_ai(user_input)
+            st.session_state.chat_history.append({"role": "assistant", "content": response})
+            st.chat_message("assistant").write(response)