Update app.py

app.py

@@ -7,22 +7,17 @@ import pdfplumber
 import difflib
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-import streamlit.components.v1 as components
 
-# ========== CONFIGURATION ==========
-st.set_page_config(
-    layout="wide",
-    page_title="Contract Analysis Suite",
-    page_icon="📑"
-)
+# ========== CONFIG ==========
+st.set_page_config(layout="wide", page_title="Contract Analysis Suite", page_icon="📑")
 
-# Initialize session state variables if they don't exist
+# ========== SESSION STATE ==========
 if 'comparison_results' not in st.session_state:
     st.session_state.comparison_results = None
 if 'analysis_results' not in st.session_state:
     st.session_state.analysis_results = None
 
-# ========== CACHED DATA LOADING ==========
+# ========== CACHED HELPERS ==========
 @st.cache_data(show_spinner=False)
 def load_questions():
     try:
@@ -41,293 +36,209 @@ def load_questions_short():
         st.error(f"Error loading short questions: {str(e)}")
         return []
 
-# ========== UTILITY FUNCTIONS ==========
+# ========== FILE PARSING ==========
 def extract_text_from_pdf(uploaded_file):
     try:
         with pdfplumber.open(uploaded_file) as pdf:
             full_text = ""
             for page in pdf.pages:
                 try:
-                    text = page.extract_text_formatted()  # Try to get formatted text
+                    text = page.extract_text_formatted()
                 except AttributeError:
                     text = page.extract_text()
-                if text:
-                    full_text += text + "\n\n"  # Add page separator
-                else:
-                    full_text += page.extract_text() + "\n\n"
-            return full_text if full_text.strip() else ""
+                full_text += (text or "") + "\n\n"
+            return full_text.strip()
     except Exception as e:
         st.error(f"PDF extraction error: {str(e)}")
         return ""
 
+def load_contract(file):
+    if not file:
+        return ""
+    try:
+        ext = file.name.split('.')[-1].lower()
+        if ext == 'txt':
+            return StringIO(file.getvalue().decode("utf-8")).read().strip()
+        elif ext == 'pdf':
+            content = extract_text_from_pdf(file)
+            if not content:
+                pdfReader = PyPDF4.PdfFileReader(file)
+                return "\n\n".join([p.extractText() for p in pdfReader.pages])
+            return content
+        elif ext == 'docx':
+            return docx2txt.process(file).strip()
+        else:
+            st.warning("Unsupported file type")
+            return ""
+    except Exception as e:
+        st.error(f"Error loading file: {str(e)}")
+        return ""
 
-
+# ========== TEXT UTILS ==========
 def highlight_differences_words(text1, text2):
     differ = difflib.Differ()
     diff = list(differ.compare(text1.split(), text2.split()))
-
-    highlighted_text1 = ""
-    highlighted_text2 = ""
-
+    h1, h2 = "", ""
     for i, word in enumerate(diff):
         if word.startswith("- "):
-            removed_word = word[2:]
-            highlighted_text1 += f'<span style="background-color:#ffcccc; display: inline-block;">{removed_word}</span>'
-            # Check for corresponding addition to highlight as changed
-            if i + 1 < len(diff) and diff[i + 1].startswith("+ "):
-                added_word = diff[i + 1][2:]
-                highlighted_text2 += f'<span style="background-color:#ffffcc; display: inline-block;">{added_word}</span>'  # Yellow for changed in text2
-                diff[i + 1] = '  '  # Consume the addition
+            w = word[2:]
+            h1 += f'<span style="background-color:#ffcccc;">{w}</span> '
+            if i+1 < len(diff) and diff[i+1].startswith("+ "):
+                h2 += f'<span style="background-color:#ffffcc;">{diff[i+1][2:]}</span> '
+                diff[i+1] = '  '
             else:
-                highlighted_text2 += " "
+                h2 += " "
         elif word.startswith("+ "):
-            added_word = word[2:]
-            highlighted_text2 += f'<span style="background-color:#ccffcc; display: inline-block;">{added_word}</span>'
-            # Check for corresponding removal
-            if i - 1 >= 0 and diff[i - 1].startswith("- "):
-                highlighted_text1 += f'<span style="background-color:#ffffcc; display: inline-block;">{diff[i-1][2:]}</span>' # Yellow for changed in text1
+            w = word[2:]
+            h2 += f'<span style="background-color:#ccffcc;">{w}</span> '
+            if i-1 >= 0 and diff[i-1].startswith("- "):
+                h1 += f'<span style="background-color:#ffffcc;">{diff[i-1][2:]}</span> '
                 diff[i-1] = '  '
             else:
-                highlighted_text1 += " "
-
+                h1 += " "
         elif word.startswith("  "):
-            highlighted_text1 += word[2:] + " "
-            highlighted_text2 += word[2:] + " "
+            w = word[2:] + " "
+            h1 += w
+            h2 += w
+    return h1.strip(), h2.strip()
 
-    return highlighted_text1, highlighted_text2
 def calculate_similarity(text1, text2):
     if not text1.strip() or not text2.strip():
         return 0.0
-
     try:
         vectorizer = TfidfVectorizer(token_pattern=r'(?u)\b\w+\b')
-        tfidf_matrix = vectorizer.fit_transform([text1, text2])
-        similarity = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:2])
-        return similarity[0][0] * 100
-    except ValueError:
+        tfidf = vectorizer.fit_transform([text1, text2])
+        sim = cosine_similarity(tfidf[0:1], tfidf[1:2])
+        return sim[0][0] * 100
+    except:
         return difflib.SequenceMatcher(None, text1, text2).ratio() * 100
 
-def load_contract(file):
-    if file is None:
-        return ""
-
-    ext = file.name.split('.')[-1].lower()
-    try:
-        if ext == 'txt':
-            content = StringIO(file.getvalue().decode("utf-8")).read()
-        elif ext == 'pdf':
-            content = extract_text_from_pdf(file)
-            if not content:
-                # Fallback to PyPDF4
-                pdfReader = PyPDF4.PdfFileReader(file)
-                full_text = ""
-                for page in pdfReader.pages:
-                    text = page.extractText()
-                    if text:
-                        full_text += text + "\n\n"
-                content = full_text
-        elif ext == 'docx':
-            content = docx2txt.process(file)
-        else:
-            st.warning('Unsupported file type')
-            return ""
-        return content.strip() if content else ""
-    except Exception as e:
-        st.error(f"Error loading {ext.upper()} file: {str(e)}")
-        return ""
-
 # ========== MAIN APP ==========
 def main():
+    st.title("📑 Contract Analysis Suite")
+    st.markdown("Compare documents and analyze legal clauses using AI-powered tools.")
+
     questions = load_questions()
     questions_short = load_questions_short()
 
     if not questions or not questions_short or len(questions) != len(questions_short):
-        st.error("Failed to load questions or questions mismatch. Please check data files.")
+        st.error("Questions failed to load properly.")
         return
 
-    st.title("📑 Contract Analysis Suite")
-    st.markdown("""
-    Compare documents and analyze legal clauses using AI-powered question answering.
-    """)
-
-    # ===== DOCUMENT UPLOAD SECTION =====
     st.header("1. Upload Documents")
     col1, col2 = st.columns(2)
 
     with col1:
-        uploaded_file1 = st.file_uploader(
-            "Upload First Document",
-            type=["txt", "pdf", "docx"],
-            key="file1"
-        )
-        contract_text1 = load_contract(uploaded_file1) if uploaded_file1 else ""
-        doc1_container = st.empty()
+        file1 = st.file_uploader("Upload First Document", type=["txt", "pdf", "docx"], key="file1")
+        text1 = load_contract(file1) if file1 else ""
+        display1 = st.empty()
 
     with col2:
-        uploaded_file2 = st.file_uploader(
-            "Upload Second Document",
-            type=["txt", "pdf", "docx"],
-            key="file2"
-        )
-        contract_text2 = load_contract(uploaded_file2) if uploaded_file2 else ""
-        doc2_container = st.empty()
+        file2 = st.file_uploader("Upload Second Document", type=["txt", "pdf", "docx"], key="file2")
+        text2 = load_contract(file2) if file2 else ""
+        display2 = st.empty()
 
-    # Update document displays with synchronized scrolling
-    if uploaded_file1:
-        doc1_content = f'<div style="border:1px solid #ccc; padding:10px; white-space: pre-wrap; font-family: monospace; font-size: 0.9em; height: 400px; overflow-y: auto;" id="doc1_text">{contract_text1}</div>'
-        doc1_container.markdown(doc1_content, unsafe_allow_html=True)
-    if uploaded_file2:
-        doc2_content = f'<div style="border:1px solid #ccc; padding:10px; white-space: pre-wrap; font-family: monospace; font-size: 0.9em; height: 400px; overflow-y: auto;" id="doc2_text">{contract_text2}</div>'
-        doc2_container.markdown(doc2_content, unsafe_allow_html=True)
+    if file1:
+        display1.text_area("Document 1 Content", value=text1, height=400, key="area1")
+    if file2:
+        display2.text_area("Document 2 Content", value=text2, height=400, key="area2")
 
-    # JavaScript for synchronized scrolling
-    scroll_script = """
-    <script>
-    function syncScroll(id, otherId) {
-        var element = document.getElementById(id);
-        var otherElement = document.getElementById(otherId);
-        if (element && otherElement) {
-            element.addEventListener('scroll', function() {
-                otherElement.scrollTop = element.scrollTop;
-            });
-            otherElement.addEventListener('scroll', function() {
-                element.scrollTop = otherElement.scrollTop;
-            });
-        }
-    }
-    window.onload = function() {
-        syncScroll('doc1_text', 'doc2_text');
-    };
-    </script>
-    """
-    components.html(scroll_script, height=0)
-
-    if not (uploaded_file1 and uploaded_file2):
-        st.warning("Please upload both documents to proceed")
+    if not (file1 and file2):
+        st.warning("Please upload both documents.")
         return
 
-    # ===== DOCUMENT COMPARISON SECTION =====
     st.header("2. Document Comparison")
-
     with st.expander("Show Document Differences", expanded=True):
         if st.button("Compare Documents"):
-            with st.spinner("Analyzing documents..."):
-                if not contract_text1.strip() or not contract_text2.strip():
-                    st.error("One or both documents appear to be empty or couldn't be read properly")
-                    return
-
-                similarity_score = calculate_similarity(contract_text1, contract_text2)
-
-
-                highlighted_diff1, highlighted_diff2 = highlight_differences_words(contract_text1, contract_text2)
+            with st.spinner("Analyzing..."):
+                sim = calculate_similarity(text1, text2)
+                diff1, diff2 = highlight_differences_words(text1, text2)
                 st.session_state.comparison_results = {
-                    'similarity_score': similarity_score,
-                    'highlighted_diff1': highlighted_diff1,
-                    'highlighted_diff2': highlighted_diff2,
-
+                    'similarity': sim,
+                    'diff1': diff1,
+                    'diff2': diff2,
                 }
 
-
-        # Display comparison results
         if st.session_state.comparison_results:
-            st.metric("Document Similarity Score",
-                        f"{st.session_state.comparison_results['similarity_score']:.2f}%")
-
-            if st.session_state.comparison_results['similarity_score'] <= 70:
-                st.warning("Significant differences detected")
-
-            st.markdown("**Visual Difference Highlighting:**")
-
-            col1_diff, col2_diff = st.columns(2)
-            with col1_diff:
-                st.markdown("### Original Document")
-                diff1_content = f'<div style="border:1px solid #ccc; padding:10px; white-space: pre-wrap; font-family: monospace; font-size: 0.9em; max-height: 500px; overflow-y: auto;" id="diff1_text">{st.session_state.comparison_results["highlighted_diff1"]}</div>'
-                st.markdown(diff1_content, unsafe_allow_html=True)
-            with col2_diff:
-                st.markdown("### Modified Document")
-                diff2_content = f'<div style="border:1px solid #ccc; padding:10px; white-space: pre-wrap; font-family: monospace; font-size: 0.9em; max-height: 500px; overflow-y: auto;" id="diff2_text">{st.session_state.comparison_results["highlighted_diff2"]}</div>'
-                st.markdown(diff2_content, unsafe_allow_html=True)
-
-            # JavaScript for synchronized scrolling of diff panes
-            diff_scroll_script = """
-            <script>
-            function syncDiffScroll(id, otherId) {
-                var element = document.getElementById(id);
-                var otherElement = document.getElementById(otherId);
-                if (element && otherElement) {
-                    element.addEventListener('scroll', function() {
-                        otherElement.scrollTop = element.scrollTop;
-                    });
-                    otherElement.addEventListener('scroll', function() {
-                        element.scrollTop = otherElement.scrollTop;
-                    });
-                }
-            }
-            window.onload = function() {
-                syncDiffScroll('diff1_text', 'diff2_text');
-            };
-            </script>
-            """
-            components.html(diff_scroll_script, height=0)
-
+            sim = st.session_state.comparison_results['similarity']
+            st.metric("Document Similarity Score", f"{sim:.2f}%")
+
+            if sim >= 70:
+                st.markdown("### Visual Difference Highlighting")
+                sync_scroll_script = """
+                <script>
+                const left = document.getElementById("left");
+                const right = document.getElementById("right");
+
+                left.onscroll = function() {
+                    right.scrollTop = left.scrollTop;
+                };
+                right.onscroll = function() {
+                    left.scrollTop = right.scrollTop;
+                };
+                </script>
+                """
+
+                html = f"""
+                <div style="display: flex; gap: 20px;">
+                    <div id="left" style="width: 100%; height: 500px; overflow-y: auto; padding: 10px; font-family: monospace; border: 1px solid #ccc;">
+                        {st.session_state.comparison_results['diff1']}
+                    </div>
+                    <div id="right" style="width: 100%; height: 500px; overflow-y: auto; padding: 10px; font-family: monospace; border: 1px solid #ccc;">
+                        {st.session_state.comparison_results['diff2']}
+                    </div>
+                </div>
+                {sync_scroll_script}
+                """
+                st.markdown(html, unsafe_allow_html=True)
+            else:
+                st.warning("Similarity below 70%. Skipping visual diff display.")
 
-    # ===== QUESTION ANALYSIS SECTION =====
+    # ========== CLAUSE ANALYSIS ==========
     st.header("3. Clause Analysis")
-
     try:
-        question_selected = st.selectbox(
-            'Select a legal question to analyze:',
-            questions_short,
-            index=0,
-            key="question_select"
-        )
-        question_idx = questions_short.index(question_selected)
-        selected_question = questions[question_idx]
-    except Exception as e:
-        st.error(f"Error selecting question: {str(e)}")
+        question_short = st.selectbox("Select a legal question to analyze:", questions_short)
+        idx = questions_short.index(question_short)
+        question = questions[idx]
+    except:
+        st.error("Error selecting question")
         return
 
     if st.button("Analyze Both Documents"):
-        if not (contract_text1.strip() and contract_text2.strip()):
-            st.error("Please ensure both documents have readable content")
+        if not (text1.strip() and text2.strip()):
+            st.error("Ensure both documents have content.")
             return
 
-        col1_analysis, col2_analysis = st.columns(2)
+        col1, col2 = st.columns(2)
 
-        with col1_analysis:
+        with col1:
             st.subheader("First Document Analysis")
-            with st.spinner('Processing first document...'):
+            with st.spinner("Processing..."):
                 try:
-                    predictions1 = run_prediction([selected_question], contract_text1, 'marshmellow77/roberta-base-cuad', n_best_size=5)
-                    answer1 = predictions1.get('0', 'No answer found')
+                    ans1 = run_prediction([question], text1, 'marshmellow77/roberta-base-cuad', n_best_size=5).get('0', 'No answer')
                     st.session_state.analysis_results = st.session_state.analysis_results or {}
-                    st.session_state.analysis_results['doc1'] = answer1 if answer1 else "No relevant clause found"
+                    st.session_state.analysis_results['doc1'] = ans1
                 except Exception as e:
-                    st.session_state.analysis_results = st.session_state.analysis_results or {}
-                    st.session_state.analysis_results['doc1'] = f"Analysis failed: {str(e)}"
+                    st.session_state.analysis_results['doc1'] = f"Failed: {e}"
 
-        with col2_analysis:
+        with col2:
             st.subheader("Second Document Analysis")
-            with st.spinner('Processing second document...'):
+            with st.spinner("Processing..."):
                 try:
-                    predictions2 = run_prediction([selected_question], contract_text2, 'marshmellow77/roberta-base-cuad', n_best_size=5)
-                    answer2 = predictions2.get('0', 'No answer found')
+                    ans2 = run_prediction([question], text2, 'marshmellow77/roberta-base-cuad', n_best_size=5).get('0', 'No answer')
                     st.session_state.analysis_results = st.session_state.analysis_results or {}
-                    st.session_state.analysis_results['doc2'] = answer2 if answer2 else "No relevant clause found"
+                    st.session_state.analysis_results['doc2'] = ans2
                 except Exception as e:
-                    st.session_state.analysis_results = st.session_state.analysis_results or {}
-                    st.session_state.analysis_results['doc2'] = f"Analysis failed: {str(e)}"
+                    st.session_state.analysis_results['doc2'] = f"Failed: {e}"
 
-    # Display analysis results
     if st.session_state.analysis_results:
-        col1_answer, col2_answer = st.columns(2)
-        with col1_answer:
-            st.subheader("First Document Analysis")
-            st.success(st.session_state.analysis_results.get('doc1', 'No analysis performed yet'))
-
-        with col2_answer:
-            st.subheader("Second Document Analysis")
-            st.success(st.session_state.analysis_results.get('doc2', 'No analysis performed yet'))
+        col1, col2 = st.columns(2)
+        with col1:
+            st.subheader("First Document Result")
+            st.success(st.session_state.analysis_results.get('doc1', 'No analysis yet'))
+        with col2:
+            st.subheader("Second Document Result")
+            st.success(st.session_state.analysis_results.get('doc2', 'No analysis yet'))
 
 if __name__ == "__main__":
-    main()
+    main()