Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,7 +12,7 @@ os.environ["HF_HOME"] = os.path.join(HOME, ".cache", "huggingface")
 os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.join(HOME, ".cache", "huggingface", "hub")
 os.environ["TRANSFORMERS_CACHE"] = os.path.join(HOME, ".cache", "huggingface", "transformers")
 os.environ["TORCH_HOME"] = os.path.join(HOME, ".cache", "torch")
-# Streamlit reads Path.home() so ~/.streamlit will resolve under our HOME
+# Create directories so first write doesn't fail
 for d in [
     HOME,
     os.environ["XDG_CACHE_HOME"],
@@ -32,8 +32,6 @@ import base64
 import hashlib
 from io import BytesIO
 from pathlib import Path
-from typing import Dict, Any
-
 import pandas as pd
 from PIL import Image
 import streamlit as st
@@ -100,11 +98,11 @@ def _get_hf_token():
         pass
     return None, None
 
-def hf_login(token: str) -> tuple[bool, str | None]:
+def hf_login(token):
     try:
         try:
             login(token, add_to_git_credential=False)
-        except TypeError:  # older hub
+        except TypeError:  # older hub versions
             login(token)
         return True, None
     except Exception as e:
@@ -128,7 +126,6 @@ if hf_token is None:
     st.session_state.logged_in = True
     st.toast("Logged in successfully.", icon="✅")
 else:
-    # Using env/secrets token; login caches to our writable HOME
     ok, err = hf_login(hf_token)
     if not ok:
         st.error(f"Failed to log in with {hf_src or 'unknown'} token: {err}")
@@ -168,11 +165,12 @@ def load_model_and_processor(hf_model_id: str, task_prompt: str):
             f"Original error: {e}"
         )
 
+    import torch
     model.eval()
-    device = __import__("torch").device("cuda" if __import__("torch").cuda.is_available() else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
 
-    with __import__("torch").no_grad():
+    with torch.no_grad():
         decoder_input_ids = processor.tokenizer(
             task_prompt,
             add_special_tokens=False,
@@ -181,7 +179,7 @@ def load_model_and_processor(hf_model_id: str, task_prompt: str):
 
     return processor, model, device, decoder_input_ids
 
-def run_inference_on_image(image: Image.Image, processor, model, device, decoder_input_ids):
+def run_inference_on_image(image, processor, model, device, decoder_input_ids):
     import torch
     pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
     gen_kwargs = dict(
@@ -211,7 +209,7 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
     return pred_dict
 
 def map_prediction_to_ui(pred):
-    import json, re
+    import re, json
     from collections import defaultdict
 
     def safe_json_load(s):
@@ -229,7 +227,8 @@ def map_prediction_to_ui(pred):
                 subs, stack, start = [], [], None
                 for i, ch in enumerate(s):
                     if ch == "{":
-                        if not stack: start = i
+                        if not stack:
+                            start = i
                         stack.append("{")
                     elif ch == "}":
                         if stack:
@@ -244,21 +243,28 @@ def map_prediction_to_ui(pred):
         return None
 
     def clean_number(x):
-        if x is None: return 0.0
-        if isinstance(x, (int, float)): return float(x)
+        if x is None:
+            return 0.0
+        if isinstance(x, (int, float)):
+            return float(x)
         s = re.sub(r"[,\s]", "", str(x).strip())
         s = re.sub(r"[^\d\.\-]", "", s)
-        if s in ("", ".", "-", "-."): return 0.0
-        try: return float(s)
-        except Exception: return 0.0
+        if s in ("", ".", "-", "-."):
+            return 0.0
+        try:
+            return float(s)
+        except Exception:
+            return 0.0
 
     def collect_keys(obj, out):
         if isinstance(obj, dict):
             for k, v in obj.items():
                 lk = str(k).strip().lower()
-                out[lk].append(v); collect_keys(v, out)
+                out[lk].append(v)
+                collect_keys(v, out)
         elif isinstance(obj, list):
-            for it in obj: collect_keys(it, out)
+            for it in obj:
+                collect_keys(it, out)
 
     def collect_lists_of_dicts(obj, out_lists):
         if isinstance(obj, dict):
@@ -275,7 +281,8 @@ def map_prediction_to_ui(pred):
                     collect_lists_of_dicts(it, out_lists)
 
     def map_item_dict(it):
-        if not isinstance(it, dict): return None
+        if not isinstance(it, dict):
+            return None
         lower = {str(k).strip().lower(): v for k, v in it.items()}
         desc = (lower.get("descriptions") or lower.get("description") or lower.get("desc") or lower.get("item") or "")
         qty = lower.get("quantity") or lower.get("qty") or lower.get("count") or ""
@@ -293,8 +300,10 @@ def map_prediction_to_ui(pred):
         }
 
     parsed = safe_json_load(pred) if isinstance(pred, str) else pred
-    if parsed is None and isinstance(pred, str): parsed = None
-    if parsed is None and not isinstance(pred, dict): parsed = pred
+    if parsed is None and isinstance(pred, str):
+        parsed = None
+    if parsed is None and not isinstance(pred, dict):
+        parsed = pred
 
     ui = {
         "Invoice Number": "", "Invoice Date": "", "Due Date": "", "Currency": "",
@@ -315,10 +324,13 @@ def map_prediction_to_ui(pred):
             lk = k.strip().lower()
             if lk in key_map:
                 for v in key_map[lk]:
-                    if v is None: continue
-                    if isinstance(v, (dict, list)): return v
+                    if v is None:
+                        continue
+                    if isinstance(v, (dict, list)):
+                        return v
                     s = str(v).strip()
-                    if s != "": return s
+                    if s != "":
+                        return s
         return None
 
     ui["Invoice Number"] = pick_first("invoice_no", "invoice_number", "invoiceid", "invoice id") or ""
@@ -333,23 +345,26 @@ def map_prediction_to_ui(pred):
     for bk in ("bank_name","bank_account_number","bank_acc_name","bank_iban","bank_swift","bank_routing","bank_branch","iban"):
         val = pick_first(bk, bk.replace("bank_", ""))
         if val:
-            bank["bank_iban" if bk=="iban" else bk] = str(val)
+            bank["bank_iban" if bk == "iban" else bk] = str(val)
     ui["Bank Details"] = bank
 
-    def _num(*keys): 
+    def _num(*keys):
         v = pick_first(*keys) or 0.0
-        try: return float(str(v).replace(",",""))
-        except Exception: return 0.0
+        try:
+            return float(str(v).replace(",", ""))
+        except Exception:
+            return 0.0
 
-    ui["Subtotal"] = _num("subtotal","sub_total","sub total")
-    ui["Tax Percentage"] = _num("tax_rate","tax_percentage","tax pct","tax percentage")
-    ui["Total Tax"] = _num("tax_amount","tax","total_tax")
-    ui["Total Amount"] = _num("total_amount","grand_total","total","amount")
+    ui["Subtotal"] = _num("subtotal", "sub_total", "sub total")
+    ui["Tax Percentage"] = _num("tax_rate", "tax_percentage", "tax pct", "tax percentage")
+    ui["Total Tax"] = _num("tax_amount", "tax", "total_tax")
+    ui["Total Amount"] = _num("total_amount", "grand_total", "total", "amount")
     ui["Currency"] = (pick_first("currency") or "").strip()
 
     items_rows = []
     def list_looks_like_items(lst):
-        if not isinstance(lst, list) or not lst or not isinstance(lst[0], dict): return False
+        if not isinstance(lst, list) or not lst or not isinstance(lst[0], dict):
+            return False
         expected = {"descriptions","description","desc","item","quantity","qty","amount","unit_price","line_total","line total"}
         return bool(expected.intersection({str(k).strip().lower() for k in lst[0].keys()}))
 
@@ -357,14 +372,17 @@ def map_prediction_to_ui(pred):
         if list_looks_like_items(cand):
             for it in cand:
                 row = map_item_dict(it)
-                if row is not None: items_rows.append(row)
-            if items_rows: break
+                if row is not None:
+                    items_rows.append(row)
+            if items_rows:
+                break
 
     if not items_rows and isinstance(parsed, dict):
         single_keys = {k.strip().lower() for k in parsed.keys()}
         if single_keys.intersection({"descriptions","description","desc","item","quantity","qty","unit_price","unit price","price","amount","line_total","line total","sku","tax","tax_amount"}):
             row = map_item_dict(parsed)
-            if row is not None: items_rows.append(row)
+            if row is not None:
+                items_rows.append(row)
 
     if not items_rows:
         for _, vals in key_map.items():
@@ -373,7 +391,8 @@ def map_prediction_to_ui(pred):
                     lower_keys = {str(x).strip().lower() for x in v.keys()}
                     if lower_keys.intersection({"descriptions","description","desc","amount","line_total","quantity","qty","unit_price"}):
                         row = map_item_dict(v)
-                        if row is not None: items_rows.append(row)
+                        if row is not None:
+                            items_rows.append(row)
 
     if not items_rows:
         desc, amt, qty, unit_price = pick_first("descriptions","description"), pick_first("amount","line_total"), pick_first("quantity","qty"), pick_first("unit_price","price")
@@ -430,16 +449,19 @@ def flatten_invoice_to_rows(invoice_data) -> list:
 
     if not line_items:
         row = base_invoice_info()
-        for k in EXPECTED_BANK_FIELDS: row[k] = bank_details.get(k, "")
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = bank_details.get(k, "")
         row.update({
             "Item Description": "", "Item Quantity": 0, "Item Unit Price": 0.0,
             "Item Amount": 0.0, "Item Tax": 0.0, "Item Line Total": 0.0,
         })
-        rows.append(row); return rows
+        rows.append(row)
+        return rows
 
     for item in line_items:
         row = base_invoice_info()
-        for k in EXPECTED_BANK_FIELDS: row[k] = bank_details.get(k, "")
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = bank_details.get(k, "")
         row.update({
             "Item Description": item.get("Description", "") if isinstance(item, dict) else "",
             "Item Quantity": item.get("Quantity", 0) if isinstance(item, dict) else 0,
@@ -641,11 +663,19 @@ elif len(st.session_state.batch_results) > 0:
                 except Exception as e:
                     st.error(f"Re-run failed: {e}")
 
-        # All editors inside one form (prevents reruns while typing)
+        # --------- All editors inside one form (prevents reruns while typing)
         with st.form(f"edit_form_{selected_hash}", clear_on_submit=False):
+            # Top bar submit (ensures Streamlit sees a submit in the form)
+            top_bar = st.columns([1, 1, 6])
+            with top_bar[0]:
+                save_top = st.form_submit_button("💾 Save", use_container_width=True)
+            with top_bar[1]:
+                # Optional no-op to satisfy UX; not used
+                st.form_submit_button("↩️ Discard (no-op)", use_container_width=True)
+
             tabs = st.tabs(["Invoice Details", "Sender/Recipient info", "Bank Details", "Line Items"])
 
-            # Invoice Details
+            # --- Invoice Details ---
             with tabs[0]:
                 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}")
@@ -662,13 +692,14 @@ elif len(st.session_state.batch_results) > 0:
                 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}")
 
-            # Sender / Recipient
+            # --- Sender / Recipient ---
             with tabs[1]:
                 sender_name = st.text_input("Sender Name*", value=form_data.get('Sender Name', ''), key=f"sender_name_{selected_hash}")
                 sender_address = st.text_area("Sender Address*", value=form_data.get('Sender Address', ''), key=f"sender_address_{selected_hash}")
                 recipient_name = st.text_input("Recipient Name*", value=form_data.get('Recipient Name', ''), key=f"recipient_name_{selected_hash}")
                 recipient_address = st.text_area("Recipient Address*", value=form_data.get('Recipient Address', ''), key=f"recipient_address_{selected_hash}")
 
+                # This button won’t submit; it flags the swap action for this run
                 if st.button("⇄ Swap", help="Swap sender and recipient information", key=f"swap_{selected_hash}"):
                     sender_name, recipient_name = recipient_name, sender_name
                     sender_address, recipient_address = recipient_address, sender_address
@@ -678,7 +709,7 @@ elif len(st.session_state.batch_results) > 0:
                 form_data['Recipient Name'] = recipient_name
                 form_data['Recipient Address'] = recipient_address
 
-            # Bank Details
+            # --- Bank Details ---
             with tabs[2]:
                 bank_details = form_data.get("Bank Details", {}) if isinstance(form_data.get("Bank Details", {}), dict) else {}
                 bank_name = st.text_input("Bank Name", value=bank_details.get('bank_name', ''), key=f"bank_name_{selected_hash}")
@@ -698,7 +729,7 @@ elif len(st.session_state.batch_results) > 0:
                 form_data["Bank Details"]["bank_routing"] = bank_routing
                 form_data["Bank Details"]["bank_branch"] = bank_branch
 
-            # Line Items (fixed height)
+            # --- Line Items (fixed height) ---
             with tabs[3]:
                 editor_key = f"item_editor_{selected_hash}"
                 item_rows = form_data.get('Itemized Data', []) or []
@@ -740,12 +771,17 @@ elif len(st.session_state.batch_results) > 0:
                     height=420,  # fixes growth/shrink reflow
                 )
 
-            # Submit (no rerun during typing)
-            submitted = st.form_submit_button("💾 Save Edits for This File")
-            if submitted:
+            # Bottom submit (near the editor)
+            save_bottom = st.form_submit_button("💾 Save Edits for This File", use_container_width=True)
+
+            # Commit once if either top or bottom button was clicked
+            if save_top or save_bottom:
+                # Make sure the table’s last edit is committed:
+                # (Press Enter or click outside the last cell before hitting Save)
                 form_data['Itemized Data'] = edited_df.to_dict('records')
                 st.session_state.batch_results[selected_hash]["edited_data"] = form_data
                 st.toast(f"Saved edits for {current['file_name']}", icon="💾")
+                st.rerun()  # one-time rerun to refresh everything
 
         # Per-file CSV download
         st.markdown("---")