Update app.py

app.py

@@ -2,9 +2,6 @@ import streamlit as st
 from PyPDF2 import PdfReader
 from anthropic import Anthropic
 from prompts import DIFFERENTIATE_PROMPT
-from concurrent.futures import ProcessPoolExecutor
-
-
 
 def extract_differences(input_text):
         input_text = input_text.strip()
@@ -48,8 +45,6 @@ def get_llm_response(extractedtext1, extractedtext2):
 
     try:
         try:
-            before_differences = message_text.split("<differences>")[0]
-            after_differences = message_text.split("</differences>")[1]
             differences_list = extract_differences(message_text.split("<differences>")[1].split("</differences>")[0].strip())
         except Exception as e:
             differences_list = extract_differences(message_text)
@@ -57,36 +52,9 @@ def get_llm_response(extractedtext1, extractedtext2):
         print("Error:", e)
         return message_text, []
     
-    difference_content = "\n\n\n".join([f"**Text1:**\n\n{d['text1']}\n\n**Text2:**\n\n{d['text2']}\n\n**Explanation:**\n\n{d['explanation']}\n\n----------------------\n" for d in differences_list])
     # display_text = f"{before_differences}\n\n{difference_content}\n\n{after_differences}"
-    display_text = difference_content
-    return display_text, differences_list
-
-def process_text_pair(pair):
-    etext1, etext2 = pair
-    dtext, pdata = get_llm_response(etext1, etext2)
-    return dtext, pdata
-
-def process_concurrently(extracted_text1, extracted_text2):
-    # Create a pool of worker processes
-    with ProcessPoolExecutor(max_workers=5) as executor:
-        # Submit tasks to the pool
-        futures = [executor.submit(process_text_pair, (etext1, etext2)) 
-                   for etext1, etext2 in zip(extracted_text1, extracted_text2)]
-        
-        # Collect results
-        display_text = ""
-        parsed_data = []
-        for future in futures:
-            result = future.result()
-            display_text += result[0]
-            parsed_data += result[1]
-    
-    # Combine results
-    # display_text = ''.join(display_texts)
-    # parsed_data = ''.join(parsed_datas)
-    
-    return display_text, parsed_data
+    return differences_list
+
 
 
 def extract_text_with_pypdf(pdf_path):
@@ -99,32 +67,39 @@ def extract_text_with_pypdf(pdf_path):
 
 def main():
     st.set_page_config(layout="wide")  # Enable wide layout
+    if "differences_data" not in  st.session_state:
+        st.session_state.differences_data = []
+    if "file1" not in st.session_state:
+        st.session_state.file1 = None
+    if "file2" not in st.session_state:
+        st.session_state.file2 = None
     st.markdown('<div style="text-align: center;">' + '<h1>PDF Upload and Compare App</h1>' + '</div>', unsafe_allow_html=True)
     
     # Create columns for side-by-side buttons
     col1, col2 = st.columns([2, 2])
         
     # Add upload button to left column
-    uploaded_file1 = col1.file_uploader("**Text 1**", type="pdf")
+    st.session_state.file1 = col1.file_uploader("**Text 1**", type="pdf")
     
     # Add upload button to right column
-    uploaded_file2 = col2.file_uploader("**Text 2**", type="pdf")
+    st.session_state.file2 = col2.file_uploader("**Text 2**", type="pdf")
+
     
     # Check if both files are uploaded
-    if uploaded_file1 and uploaded_file2:
+    if st.session_state.file1 and st.session_state.file2:
         # Get filenames from uploaded files
-        filename1 = uploaded_file1.name
-        filename2 = uploaded_file2.name
+        filename1 = st.session_state.file1.name
+        filename2 = st.session_state.file2.name
         
         try:
-            extracted_text1 = extract_text_with_pypdf(uploaded_file1)
-            extracted_text2 = extract_text_with_pypdf(uploaded_file2)
+            extracted_text1 = extract_text_with_pypdf(st.session_state.file1)
+            extracted_text2 = extract_text_with_pypdf(st.session_state.file2)
             with col1.expander(filename1):
                 st.write("\n\n".join("\n\n------------------------\n\n".join(extracted_text1).splitlines()))
             
             with col2.expander(filename2):
                 st.write("\n\n".join("\n\n------------------------\n\n".join(extracted_text2).splitlines()))
-                
+            
             st.success(f"Content of files **{filename1}** and **{filename2}** have been extracted successfully.")
         except Exception as e:
             st.error(f"Error saving files: {str(e)}")
@@ -132,32 +107,42 @@ def main():
         # Add button at the bottom to run Find Differences function
         if st.button("Find Differences"):
             try:
-                display_text = ""
-                parsed_data = []
-                # for etext1, etext2 in zip(extracted_text1, extracted_text2):
-                #     dtext, pdata = get_llm_response(etext1, etext2)
-                #     display_text += dtext
-                #     parsed_data += pdata
-                display_text, parsed_data = process_concurrently(extracted_text1, extracted_text2)
+                # display_text, parsed_data = process_concurrently(extracted_text1, extracted_text2)
                 # display_text, parsed_data = get_llm_response(extracted_text1, extracted_text2)
-                display_text1 = "\n\n------------------------\n\n".join(extracted_text1)
-                display_text2 = "\n\n------------------------\n\n".join(extracted_text2)
-
-                for diff in parsed_data:
-                    diff_text1 = diff['text1'].strip()
-                    diff_text2 = diff['text2'].strip()          
-                    if diff_text1 == "" or diff_text2 == "": continue              
-                    diff_text1_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text1.splitlines()])
-                    diff_text2_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text2.splitlines()])
-                    display_text1 = diff_text1_phrase.join(display_text1.split(diff_text1)) if diff_text1 in display_text1 else display_text1
-                    display_text2 = diff_text2_phrase.join(display_text2.split(diff_text2)) if diff_text2 in display_text2 else display_text2
-                with col1.expander(f"{filename1} Difference Highlighted"):
-                    st.write("\n\n".join(display_text1.splitlines()), unsafe_allow_html=True)
-
-                with col2.expander(f"{filename2} Difference Highlighted"):
-                    st.write("\n\n".join(display_text2.splitlines()), unsafe_allow_html=True)
+                
+                i = 1
+                for etext1, etext2 in zip(extracted_text1, extracted_text2):
+                    data = next((d for d in st.session_state.differences_data if d['etext1'] == etext1 and d['etext2'] == etext2), None)
+                    if data:
+                        pdata = data['pdata']
+                        dext1 = data['dext1']
+                        dext2 = data['dext2']
+                    else:
+                        pdata = get_llm_response(etext1, etext2)
+                        dext1 = etext1
+                        dext2 = etext2
+
+                        for diff in pdata:
+                            diff_text1 = diff['text1'].strip()
+                            diff_text2 = diff['text2'].strip()          
+                            if diff_text1 == "" or diff_text2 == "": continue              
+                            diff_text1_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text1.splitlines()])
+                            diff_text2_phrase = "\n".join([f"<span style='background-color: grey;'>{t}</span>" for t in diff_text2.splitlines()])
+                            dext1 = diff_text1_phrase.join(dext1.split(diff_text1)) if diff_text1 in dext1 else dext1
+                            dext2 = diff_text2_phrase.join(dext2.split(diff_text2)) if diff_text2 in dext2 else dext2
+
+                        st.session_state.differences_data.append({"etext1": etext1, "etext2": etext2, "pdata": pdata, "dext1": dext1, "dext2": dext2})
+                        reverse_pdata = [{'text1': d['text2'], 'text2': d['text1'], 'explanation': d['explanation']} for d in pdata]
+                        st.session_state.differences_data.append({"etext1": etext2, "etext2": etext1, "pdata": reverse_pdata, "dext1": dext2, "dext2": dext1})
                     
-                st.markdown(display_text)
+                    display_text = "\n\n\n".join([f"**Text1:**\n\n{d['text1']}\n\n**Text2:**\n\n{d['text2']}\n\n**Explanation:**\n\n{d['explanation']}\n\n----------------------\n" for d in pdata])
+
+                    with st.expander(f"**Page {i}** - {filename1}"):
+                        st.markdown("\n\n".join(dext1.splitlines()), unsafe_allow_html=True)
+                    with st.expander(f"**Page {i}** - {filename2}"):
+                        st.markdown("\n\n".join(dext2.splitlines()), unsafe_allow_html=True)
+                    st.markdown(display_text)
+                    i += 1
             except Exception as e:
                 st.error(f"Error finding differences: {str(e)}")