Update src/streamlit_app.py

src/streamlit_app.py

@@ -38,11 +38,6 @@ import pandas as pd
 from PIL import Image
 from huggingface_hub import login
 
-# 👇 ADD CALLBACK FUNCTION AT TOP LEVEL
-def trigger_rerun():
-    # Forces immediate rerun after any edit — ensures edited_df is fresh
-    pass
-
 # ---------------------------
 # UI: main
 # ---------------------------
@@ -178,7 +173,7 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
         pixel_values=pixel_values,
         decoder_input_ids=decoder_input_ids,
         max_length=1536,
-        num_beams=1,
+        num_beams=4,
         early_stopping=False,
     )
 
@@ -681,15 +676,15 @@ if not st.session_state.is_processing_batch and len(st.session_state.batch_resul
 # RESULTS VIEW — Show selector + editable form
 # ---------------------------
 elif len(st.session_state.batch_results) > 0:
-            # ---------------------------
-        # Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
-        # ---------------------------
+    # ---------------------------
+    # Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
+    # ---------------------------
     if st.button("📦 Download All Results (Excel)", key="download_all"):
-            # Collect rows from all invoices and concatenate into one DataFrame
+        # Collect rows from all invoices and concatenate into one DataFrame
         all_rows = []
         for file_hash, result in st.session_state.batch_results.items():
             rows = flatten_invoice_to_rows(result["edited_data"])
-                # Annotate rows with source file name so user can identify which invoice each row came from
+            # Annotate rows with source file name so user can identify which invoice each row came from
             for r in rows:
                 r["Source File"] = result.get("file_name", file_hash)
             all_rows.extend(rows)
@@ -699,13 +694,13 @@ elif len(st.session_state.batch_results) > 0:
         else:
             full_df = pd.DataFrame(all_rows)
 
-                # Reorder columns to put Source File first
+            # Reorder columns to put Source File first
             cols = list(full_df.columns)
             if "Source File" in cols:
                 cols = ["Source File"] + [c for c in cols if c != "Source File"]
             full_df = full_df[cols]
 
-                # Try to write XLSX (preferred). If engine not available, fall back to CSV.
+            # Try to write XLSX (preferred). If engine not available, fall back to CSV.
             buffer = BytesIO()
             dl_filename = "all_extracted_invoices.xlsx"
             tried_xlsx = False
@@ -726,7 +721,7 @@ elif len(st.session_state.batch_results) > 0:
                 dl_filename = "all_extracted_invoices.csv"
                 mime = "text/csv"
 
-                # Trigger immediate download via a data URI and small HTML snippet
+            # Trigger immediate download via a data URI and small HTML snippet
             import base64
             import streamlit.components.v1 as components
             b64 = base64.b64encode(file_bytes).decode()
@@ -768,8 +763,9 @@ elif len(st.session_state.batch_results) > 0:
     # Get current file data
     current = st.session_state.batch_results[selected_hash]
     image = current["image"]
-    import copy
-    form_data = copy.deepcopy(current["edited_data"]) 
+    
+    # ✅ FIX: Don't create a copy here - just reference the stored data
+    form_data = current["edited_data"]
 
     # Layout
     left_col, right_col = st.columns([1, 1])
@@ -806,6 +802,7 @@ elif len(st.session_state.batch_results) > 0:
                     st.rerun()
                 except Exception as e:
                     st.error(f"Re-run failed: {e}")
+                    
         tabs = st.tabs(["Invoice Details", "Sender/Recipient info", "Bank Details", "Line Items"])
 
         st.markdown(
@@ -830,48 +827,53 @@ elif len(st.session_state.batch_results) > 0:
             """,
             unsafe_allow_html=True,
         )
+
         # ---------- Invoice Details ----------
+        # ✅ FIX: Read values directly from widgets without assigning back to form_data
         with tabs[0]:
             with st.container():
-                form_data['Invoice Number'] = st.text_input("Invoice Number", value=form_data.get('Invoice Number', ''), key=f"invoice_number_{selected_hash}")
-                form_data['Invoice Date'] = st.text_input("Invoice Date", value=str(form_data.get('Invoice Date', '')).strip(), key=f"invoice_date_text_{selected_hash}")
-                form_data['Due Date'] = st.text_input("Due Date", value=str(form_data.get('Due Date', '')).strip(), key=f"due_date_text_{selected_hash}")
+                st.text_input("Invoice Number", value=form_data.get('Invoice Number', ''), key=f"invoice_number_{selected_hash}")
+                st.text_input("Invoice Date", value=str(form_data.get('Invoice Date', '')).strip(), key=f"invoice_date_text_{selected_hash}")
+                st.text_input("Due Date", value=str(form_data.get('Due Date', '')).strip(), key=f"due_date_text_{selected_hash}")
+                
                 curr_options = ['USD', 'EUR', 'GBP', 'INR', 'Other']
                 curr_value = form_data.get('Currency', 'USD')
                 curr_index = curr_options.index(curr_value) if curr_value in curr_options else (len(curr_options) - 1)
-                new_curr = st.selectbox("Currency", options=curr_options, index=curr_index, key=f"currency_select_{selected_hash}")
-                if new_curr == 'Other':
-                    new_curr = st.text_input("Specify Currency", value=form_data.get('Currency', ''), key=f"custom_currency_{selected_hash}")
-                form_data['Currency'] = new_curr
-                form_data['Subtotal'] = safe_number_input("Subtotal", form_data.get('Subtotal', 0.0), f"subtotal_{selected_hash}")
-                form_data['Tax Percentage'] = safe_number_input("Tax Percentage", form_data.get('Tax Percentage', 0.0), f"tax_pct_{selected_hash}")
-                form_data['Total Tax'] = safe_number_input("Total Tax", form_data.get('Total Tax', 0.0), f"total_tax_{selected_hash}")
-                form_data['Total Amount'] = safe_number_input("Total Amount", form_data.get('Total Amount', 0.0), f"total_amount_{selected_hash}")
+                st.selectbox("Currency", options=curr_options, index=curr_index, key=f"currency_select_{selected_hash}")
+                if st.session_state.get(f"currency_select_{selected_hash}") == 'Other':
+                    st.text_input("Specify Currency", value=form_data.get('Currency', ''), key=f"custom_currency_{selected_hash}")
+                
+                safe_number_input("Subtotal", form_data.get('Subtotal', 0.0), f"subtotal_{selected_hash}")
+                safe_number_input("Tax Percentage", form_data.get('Tax Percentage', 0.0), f"tax_pct_{selected_hash}")
+                safe_number_input("Total Tax", form_data.get('Total Tax', 0.0), f"total_tax_{selected_hash}")
+                safe_number_input("Total Amount", form_data.get('Total Amount', 0.0), f"total_amount_{selected_hash}")
 
         # ---------- Sender / Recipient ----------
         with tabs[1]:
-            # Get FLAT top-level keys (NOT nested dictionaries)
             sender_name = form_data.get('Sender Name', '')
             sender_address = form_data.get('Sender Address', '')
             recipient_name = form_data.get('Recipient Name', '')
             recipient_address = form_data.get('Recipient Address', '')
             
             with st.container():
-                sender_name = st.text_input("Sender Name*", value=sender_name, key=f"sender_name_{selected_hash}")
-                sender_address = st.text_area("Sender Address*", value=sender_address, key=f"sender_address_{selected_hash}")
-                recipient_name = st.text_input("Recipient Name*", value=recipient_name, key=f"recipient_name_{selected_hash}")
-                recipient_address = st.text_area("Recipient Address*", value=recipient_address, key=f"recipient_address_{selected_hash}")
+                st.text_input("Sender Name*", value=sender_name, key=f"sender_name_{selected_hash}")
+                st.text_area("Sender Address*", value=sender_address, key=f"sender_address_{selected_hash}")
+                st.text_input("Recipient Name*", value=recipient_name, key=f"recipient_name_{selected_hash}")
+                st.text_area("Recipient Address*", value=recipient_address, key=f"recipient_address_{selected_hash}")
                 
                 if st.button("⇄ Swap", help="Swap sender and recipient information", key=f"swap_{selected_hash}"):
-                    # Swap individual fields (NOT nested objects)
-                    form_data['Sender Name'], form_data['Recipient Name'] = form_data['Recipient Name'], form_data['Sender Name']
-                    form_data['Sender Address'], form_data['Recipient Address'] = form_data['Recipient Address'], form_data['Sender Address']
+                    # Swap in session_state widget values
+                    temp_name = st.session_state.get(f"sender_name_{selected_hash}", "")
+                    temp_addr = st.session_state.get(f"sender_address_{selected_hash}", "")
+                    
+                    st.session_state[f"sender_name_{selected_hash}"] = st.session_state.get(f"recipient_name_{selected_hash}", "")
+                    st.session_state[f"sender_address_{selected_hash}"] = st.session_state.get(f"recipient_address_{selected_hash}", "")
+                    st.session_state[f"recipient_name_{selected_hash}"] = temp_name
+                    st.session_state[f"recipient_address_{selected_hash}"] = temp_addr
                     st.rerun()
 
-        # ---------- Bank Details ---------- (FIXED)
-        # ---------- Bank Details ---------- (FIXED)
+        # ---------- Bank Details ----------
         with tabs[2]:
-            # Get bank details from nested dictionary
             bank_details = form_data.get("Bank Details", {})
             if not isinstance(bank_details, dict):
                 bank_details = {}
@@ -885,31 +887,20 @@ elif len(st.session_state.batch_results) > 0:
             bank_branch = bank_details.get('bank_branch', '')
             
             with st.container():
-                bank_name = st.text_input("Bank Name", value=bank_name, key=f"bank_name_{selected_hash}")
-                bank_acc_no = st.text_input("Account Number", value=bank_acc_no, key=f"bank_acc_no_{selected_hash}")
-                bank_acc_name = st.text_input("Bank Account Name", value=bank_acc_name, key=f"bank_acc_name_{selected_hash}")
-                bank_iban = st.text_input("IBAN", value=bank_iban, key=f"iban_{selected_hash}")
-                bank_swift = st.text_input("SWIFT Code", value=bank_swift, key=f"swift_code_{selected_hash}")
-                bank_routing = st.text_input("Routing Number", value=bank_routing, key=f"routing_{selected_hash}")
-                bank_branch = st.text_input("Branch", value=bank_branch, key=f"branch_{selected_hash}")
-            
-            # Update the nested Bank Details dictionary
-            form_data.setdefault("Bank Details", {})
-            form_data["Bank Details"]["bank_name"] = bank_name
-            form_data["Bank Details"]["bank_acc_no"] = bank_acc_no
-            form_data["Bank Details"]["bank_acc_name"] = bank_acc_name
-            form_data["Bank Details"]["bank_iban"] = bank_iban
-            form_data["Bank Details"]["bank_swift"] = bank_swift
-            form_data["Bank Details"]["bank_routing"] = bank_routing
-            form_data["Bank Details"]["bank_branch"] = bank_branch
+                st.text_input("Bank Name", value=bank_name, key=f"bank_name_{selected_hash}")
+                st.text_input("Account Number", value=bank_acc_no, key=f"bank_acc_no_{selected_hash}")
+                st.text_input("Bank Account Name", value=bank_acc_name, key=f"bank_acc_name_{selected_hash}")
+                st.text_input("IBAN", value=bank_iban, key=f"iban_{selected_hash}")
+                st.text_input("SWIFT Code", value=bank_swift, key=f"swift_code_{selected_hash}")
+                st.text_input("Routing Number", value=bank_routing, key=f"routing_{selected_hash}")
+                st.text_input("Branch", value=bank_branch, key=f"branch_{selected_hash}")
 
         # ---------- Line Items ----------
-                # ---------- Line Items ----------
         with tabs[3]:
             editor_key = f"item_editor_{selected_hash}"
             item_rows = form_data.get('Itemized Data', []) or []
 
-            # --- Normalize item keys produced by the model (e.g. "Item Description") into simple keys ---
+            # --- Normalize item keys produced by the model ---
             def normalize_item_keys(item):
                 if not isinstance(item, dict):
                     return {
@@ -938,17 +929,13 @@ elif len(st.session_state.batch_results) > 0:
                     "line_total": "Line Total",
                 }
                 new = {}
-                # map keys
                 for k, v in item.items():
                     key = mapping.get(k, mapping.get(str(k).lower(), k))
-                    # normalize to our canonical set where possible
                     if key in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                         new[key] = v
                     else:
-                        # store unknowns too so user can see them if needed
                         new[k] = v
 
-                # ensure canonical columns exist (avoid missing columns in DataFrame)
                 for kk in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                     if kk not in new:
                         new[kk] = ""
@@ -956,11 +943,8 @@ elif len(st.session_state.batch_results) > 0:
                 return new
 
             normalized_items = [normalize_item_keys(it) for it in item_rows]
-
-            # Create DataFrame from normalized items (these columns will be visible)
             df = pd.DataFrame(normalized_items)
 
-            # Ensure columns exist and are named user-friendly
             for col in ["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]:
                 if col not in df.columns:
                     df[col] = ""
@@ -975,38 +959,57 @@ elif len(st.session_state.batch_results) > 0:
             if len(edited_df) == 0:
                 st.info("No line items found in the invoice.")
 
-
-        # Save button (per file)
+        # ✅ FIX: Save button now collects values from session_state widgets
         if st.button("💾 Save Edits for This File", key=f"save_{selected_hash}"):
-            # Update line items from data editor
-            form_data['Itemized Data'] = edited_df.to_dict('records')
-            # Update session state with complete form data
-            st.session_state.batch_results[selected_hash]["edited_data"] = form_data
+            # Collect all values from session_state
+            updated_data = {
+                'Invoice Number': st.session_state.get(f"invoice_number_{selected_hash}", ""),
+                'Invoice Date': st.session_state.get(f"invoice_date_text_{selected_hash}", ""),
+                'Due Date': st.session_state.get(f"due_date_text_{selected_hash}", ""),
+                'Currency': st.session_state.get(f"custom_currency_{selected_hash}", "") if st.session_state.get(f"currency_select_{selected_hash}") == 'Other' else st.session_state.get(f"currency_select_{selected_hash}", "USD"),
+                'Subtotal': st.session_state.get(f"subtotal_{selected_hash}", 0.0),
+                'Tax Percentage': st.session_state.get(f"tax_pct_{selected_hash}", 0.0),
+                'Total Tax': st.session_state.get(f"total_tax_{selected_hash}", 0.0),
+                'Total Amount': st.session_state.get(f"total_amount_{selected_hash}", 0.0),
+                'Sender Name': st.session_state.get(f"sender_name_{selected_hash}", ""),
+                'Sender Address': st.session_state.get(f"sender_address_{selected_hash}", ""),
+                'Recipient Name': st.session_state.get(f"recipient_name_{selected_hash}", ""),
+                'Recipient Address': st.session_state.get(f"recipient_address_{selected_hash}", ""),
+                'Bank Details': {
+                    'bank_name': st.session_state.get(f"bank_name_{selected_hash}", ""),
+                    'bank_acc_no': st.session_state.get(f"bank_acc_no_{selected_hash}", ""),
+                    'bank_acc_name': st.session_state.get(f"bank_acc_name_{selected_hash}", ""),
+                    'bank_iban': st.session_state.get(f"iban_{selected_hash}", ""),
+                    'bank_swift': st.session_state.get(f"swift_code_{selected_hash}", ""),
+                    'bank_routing': st.session_state.get(f"routing_{selected_hash}", ""),
+                    'bank_branch': st.session_state.get(f"branch_{selected_hash}", "")
+                },
+                'Itemized Data': edited_df.to_dict('records')
+            }
+            
+            # Also set convenience fields
+            updated_data['Sender'] = {"Name": updated_data['Sender Name'], "Address": updated_data['Sender Address']}
+            updated_data['Recipient'] = {"Name": updated_data['Recipient Name'], "Address": updated_data['Recipient Address']}
+            
+            # Update session state
+            st.session_state.batch_results[selected_hash]["edited_data"] = updated_data
             st.success(f"✅ Edits saved for {current['file_name']}")
 
         # Download buttons (per file)
         st.markdown("---")
         col_a, col_b, col_c = st.columns([1, 1, 1])
-        #with col_a:
-            #jsonl_str = json.dumps(data, ensure_ascii=False, indent=2)
-            #st.download_button(
-               # "📥 Download JSON",
-                #jsonl_str.encode("utf-8"),
-                #file_name=f"{Path(current['file_name']).stem}_extracted.json",
-                #mime="application/json",
-                #key=f"dl_json_{selected_hash}"
-            #)
+        
         with col_b:
-            # ✅ Flatten entire invoice into rows (one per line item)
-            rows = flatten_invoice_to_rows(form_data)
+            # Use the saved edited_data (not the temporary form_data)
+            rows = flatten_invoice_to_rows(st.session_state.batch_results[selected_hash]["edited_data"])
             full_df = pd.DataFrame(rows)
             
             # Optional: Reorder columns for better readability
             desired_col_order = [
                 "Invoice Number", "Invoice Date", "Due Date", "Currency",
-                "Sender Name", "Sender Address", "Recipient Name", "Recipient Address",
                 "Subtotal", "Tax Percentage", "Total Tax", "Total Amount",
-                "bank_name", "bank_acc_no", "bank_iban", "bank_swift", "bank_routing", "bank_branch", "bank_acc_name",
+                "Sender Name", "Sender Address", "Recipient Name", "Recipient Address",
+                "bank_name", "bank_acc_no", "bank_acc_name", "bank_iban", "bank_swift", "bank_routing", "bank_branch", 
                 "Item Description", "Item Quantity", "Item Unit Price", "Item Amount", "Item Tax", "Item Line Total"
             ]
             # Keep only columns that exist
@@ -1025,11 +1028,6 @@ elif len(st.session_state.batch_results) > 0:
                 mime="text/csv",
                 key=f"dl_csv_{selected_hash}"
             )
-# Global Download All — produce a single Excel file (concatenated rows) and trigger direct download
-
-
-# ---------------------------
-# PROCESSING STATE
 
 # ---------------------------
 # PROCESSING STATE — Show progress