Update src/streamlit_app.py

src/streamlit_app.py

@@ -207,25 +207,44 @@ def run_inference_on_image(image: Image.Image, processor, model, device, decoder
 # ---------------------------
 def map_prediction_to_ui(pred):
     import json, re
-    from datetime import datetime
+    from collections import defaultdict
 
+    # --- parse raw string payloads that embed JSON ---
     def safe_json_load(s):
         if s is None:
             return None
         if isinstance(s, (dict, list)):
             return s
         if isinstance(s, str):
+            s = s.strip()
+            if s == "":
+                return None
             try:
                 return json.loads(s)
             except Exception:
-                try:
-                    t = s.strip()
-                    t = t.replace("\\'", "'").replace('\"{', '{').replace('}\"', '}')
-                    return json.loads(t)
-                except Exception:
-                    return None
+                # try to extract balanced-brace substrings (simple approach)
+                subs = []
+                stack = []
+                start = None
+                for i, ch in enumerate(s):
+                    if ch == "{":
+                        if not stack:
+                            start = i
+                        stack.append("{")
+                    elif ch == "}":
+                        if stack:
+                            stack.pop()
+                            if not stack and start is not None:
+                                subs.append(s[start:i+1])
+                                start = None
+                for sub in subs:
+                    try:
+                        return json.loads(sub)
+                    except Exception:
+                        continue
         return None
 
+    # --- normalize numeric strings like "1,800.00" -> float ---
     def clean_number(x):
         if x is None:
             return 0.0
@@ -234,6 +253,7 @@ def map_prediction_to_ui(pred):
         s = str(x).strip()
         if s == "":
             return 0.0
+        # remove commas and non-number chars except dot and minus
         s = re.sub(r"[,\s]", "", s)
         s = re.sub(r"[^\d\.\-]", "", s)
         if s in ("", ".", "-", "-."):
@@ -243,28 +263,69 @@ def map_prediction_to_ui(pred):
         except Exception:
             return 0.0
 
-    def parse_date(s):
-        if not s:
-            return ""
-        s = str(s).strip()
-        for fmt in ("%Y-%m-%d", "%d-%m-%Y", "%d/%m/%Y", "%m/%d/%Y", "%d.%m.%Y"):
-            try:
-                return datetime.strptime(s, fmt).strftime("%Y-%m-%d")
-            except Exception:
-                pass
-        m = re.match(r"^(\d{1,2})/(\d{1,2})/(\d{4})$", s)
-        if m:
-            a, b, y = int(m.group(1)), int(m.group(2)), int(m.group(3))
-            if a > 12:
-                d, mo = a, b
-            else:
-                mo, d = a, b
-            try:
-                return datetime(year=y, month=mo, day=d).strftime("%Y-%m-%d")
-            except Exception:
-                return s
-        return s
+    # --- collect all keys -> list of values from pred, recursively ---
+    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)
+        elif isinstance(obj, list):
+            for it in obj:
+                collect_keys(it, out)
+        else:
+            # primitive: handled via parent key
+            pass
+
+    # --- find list-of-dicts candidates for items (recursively) ---
+    def collect_lists_of_dicts(obj, out_lists):
+        if isinstance(obj, dict):
+            for v in obj.values():
+                if isinstance(v, list) and v and isinstance(v[0], dict):
+                    out_lists.append(v)
+                else:
+                    collect_lists_of_dicts(v, out_lists)
+        elif isinstance(obj, list):
+            for it in obj:
+                if isinstance(it, list) and it and isinstance(it[0], dict):
+                    out_lists.append(it)
+                else:
+                    collect_lists_of_dicts(it, out_lists)
+
+    # --- map item dict -> UI item row using the keys you specified in example ---
+    def map_item_dict(it):
+        if not isinstance(it, dict):
+            return None
+        # lowered keys mapping
+        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 ""
+        unit_price = lower.get("unit_price") or lower.get("price") or ""
+        amount = lower.get("amount") or lower.get("line_total") or lower.get("line total") or lower.get("total") or ""
+        tax = lower.get("tax") or lower.get("tax_amount") or ""
+        line_total = lower.get("line_total") or lower.get("line_total".lower()) or lower.get("line total") or amount
+
+        return {
+            "Description": str(desc).strip(),
+            "Quantity": float(clean_number(qty)),
+            "Unit Price": float(clean_number(unit_price)),
+            "Amount": float(clean_number(amount)),
+            "Tax": float(clean_number(tax)),
+            "Line Total": float(clean_number(line_total))
+        }
+
+    # ----------------- Start mapping -----------------
+    # Try parse if pred is a JSON-like string
+    parsed = safe_json_load(pred) if isinstance(pred, str) else pred
+    if parsed is None and isinstance(pred, str):
+        # not parseable -> fallback to empty UI
+        parsed = None
+
+    if parsed is None and not isinstance(pred, dict):
+        # nothing we can map
+        parsed = pred  # will still allow collect_keys if it's dict; else produce empty ui
 
+    # create empty UI template
     ui = {
         "Invoice Number": "",
         "Invoice Date": "",
@@ -284,137 +345,135 @@ def map_prediction_to_ui(pred):
         "Itemized Data": []
     }
 
-    if pred is None:
-        return ui
-
-    if isinstance(pred, str):
-        parsed = safe_json_load(pred)
-        if parsed is not None:
-            pred = parsed
-
-    gt = None
-    if isinstance(pred, dict):
-        if "gt_parse" in pred:
-            gp = pred["gt_parse"]
-            gp_parsed = safe_json_load(gp)
-            gt = gp_parsed if gp_parsed is not None else (gp if isinstance(gp, dict) else {})
-        else:
-            gt = pred
-    else:
-        return ui
-
-    header = gt.get("header") or {}
-    items = gt.get("items") or []
-    summary = gt.get("summary") or {}
-
-    ui["Invoice Number"] = header.get("invoice_no") or header.get("invoice_number") or ui["Invoice Number"]
-    ui["Invoice Date"] = str(header.get("invoice_date") or header.get("inv_date") or "")
-    ui["Due Date"] = str(header.get("due_date") or header.get("due") or "")
-
-    ui["Sender Name"] = header.get("sender_name") or header.get("seller_name") or header.get("from_name") or ui["Sender Name"]
-    ui["Sender Address"] = header.get("sender_addr") or header.get("sender_address") or header.get("seller_addr") or ui["Sender Address"]
-    ui["Recipient Name"] = header.get("rcpt_name") or header.get("recipient_name") or header.get("to_name") or ui["Recipient Name"]
-    ui["Recipient Address"] = header.get("rcpt_addr") or header.get("rcpt_address") or header.get("recipient_address") or ui["Recipient Address"]
+    # If parsed is a dict, collect all keys and list-of-dict candidates
+    key_map = defaultdict(list)  # lowercase-key -> list of values
+    list_candidates = []         # list of list-of-dicts found
+    if isinstance(parsed, dict):
+        collect_keys(parsed, key_map)
+        collect_lists_of_dicts(parsed, list_candidates)
+    elif isinstance(pred, dict):
+        # if parsing failed but original pred is dict, use that
+        collect_keys(pred, key_map)
+        collect_lists_of_dicts(pred, list_candidates)
+
+    # Helper to pick first non-empty value from candidate keys
+    def pick_first(*candidate_keys):
+        for k in candidate_keys:
+            lk = k.strip().lower()
+            if lk in key_map:
+                # pick first non-empty
+                for v in key_map[lk]:
+                    if v is None:
+                        continue
+                    # return primitive or string immediately; if dict/list, return as-is
+                    if isinstance(v, (dict, list)):
+                        return v
+                    s = str(v).strip()
+                    if s != "":
+                        return s
+        return None
 
-    ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
-    ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
+    # Map simple scalar fields using the exact keys you provided (plus common close variants)
+    ui["Invoice Number"] = pick_first("invoice_no", "invoice_number", "invoiceid", "invoice id") or ""
+    ui["Invoice Date"] = pick_first("invoice_date", "date", "invoice date") or ""
+    ui["Due Date"] = pick_first("due_date", "due_date", "due") or ""
+    ui["Sender Name"] = pick_first("sender_name", "sender") or ""
+    ui["Sender Address"] = pick_first("sender_addr", "sender_address", "sender addr") or ""
+    ui["Recipient Name"] = pick_first("rcpt_name", "recipient_name", "recipient", "rcpt") or ""
+    ui["Recipient Address"] = pick_first("rcpt_addr", "recipient_address", "recipient addr") or ""
 
+    # bank details: gather keys that start with 'bank_' or exact matches
     bank = {}
-    if header.get("bank_name"):
-        bank["bank_name"] = str(header.get("bank_name")).strip()
-    if header.get("bank_acc_no"):
-        bank["bank_account_number"] = str(header.get("bank_acc_no")).strip()
-    if header.get("bank_account_number"):
-        bank["bank_account_number"] = bank.get("bank_account_number") or str(header.get("bank_account_number")).strip()
-    if header.get("bank_iban"):
-        bank["bank_iban"] = str(header.get("bank_iban")).strip()
-    if header.get("bank_routing"):
-        bank["bank_routing"] = str(header.get("bank_routing")).strip()
-    if header.get("bank_swift"):
-        bank["bank_swift"] = str(header.get("bank_swift")).strip()
-    if header.get("bank_branch"):
-        bank["bank_branch"] = str(header.get("bank_branch")).strip()
-    if header.get("bank_acc_name"):
-        bank["bank_acc_name"] = str(header.get("bank_acc_name")).strip()
-    hb = header.get("bank")
-    if isinstance(hb, dict):
-        for k, v in hb.items():
-            if not v:
-                continue
-            lk = k.lower()
-            if "iban" in lk:
-                bank["bank_iban"] = bank.get("bank_iban") or str(v).strip()
-            elif "swift" in lk:
-                bank["bank_swift"] = bank.get("bank_swift") or str(v).strip()
-            elif "acc" in lk or "account" in lk:
-                bank["bank_account_number"] = bank.get("bank_account_number") or str(v).strip()
-            elif "name" in lk and "bank" in lk:
-                bank["bank_name"] = bank.get("bank_name") or str(v).strip()
-            elif "branch" in lk:
-                bank["bank_branch"] = bank.get("bank_branch") or str(v).strip()
-            elif "acc_name" in lk or "account_name" in lk:
-                bank["bank_acc_name"] = bank.get("bank_acc_name") or str(v).strip()
-
+    for bk in ("bank_name", "bank_acc_no", "bank_account_number", "bank_acc_name", "bank_iban", "bank_swift", "bank_routing", "bank_branch", "iban"):
+        val = pick_first(bk, bk.replace("bank_", ""))  # allow both 'iban' and 'bank_iban'
+        if val:
+            # normalize key name to bank_* form
+            if bk == "iban":
+                bank["bank_iban"] = str(val)
+            else:
+                bank[bk] = str(val)
     ui["Bank Details"] = bank
 
-    ui["Subtotal"] = clean_number(summary.get("subtotal") or summary.get("sub_total") or summary.get("subTotal"))
-    ui["Tax Percentage"] = clean_number(summary.get("tax_rate") or summary.get("taxRate") or summary.get("tax_percentage"))
-    ui["Total Tax"] = clean_number(summary.get("tax_amount") or summary.get("tax") or summary.get("taxAmount"))
-    ui["Total Amount"] = clean_number(summary.get("total_amount") or summary.get("grand_total") or summary.get("total") or summary.get("amount_total"))
-    ui["Currency"] = summary.get("currency") or header.get("currency") or ui["Currency"] or ""
-
-    normalized_items = []
-
-    if isinstance(items, str):
-        parsed_items = safe_json_load(items)
-        if parsed_items is not None:
-            items = parsed_items
-
-    if isinstance(items, dict):
-        if any(isinstance(v, list) for v in items.values()):
-            list_cols = {k: v for k, v in items.items() if isinstance(v, list)}
-            max_len = max((len(v) for v in list_cols.values()), default=0)
-            for i in range(max_len):
-                row = {}
-                for k, v in items.items():
-                    if isinstance(v, list):
-                        row[k] = v[i] if i < len(v) else ""
-                    else:
-                        row[k] = v
-                normalized_items.append(row)
-        else:
-            normalized_items.append(items)
-    elif isinstance(items, list):
-        normalized_items = items
-    else:
-        normalized_items = []
-
-    item_rows = []
-    for it in normalized_items:
-        if not isinstance(it, dict):
-            item_rows.append({"Description": str(it), "Quantity": 1, "Unit Price": 0.0, "Amount": 0.0, "Tax": 0.0, "Line Total": 0.0})
-            continue
-        desc = it.get("descriptions") or it.get("description") or it.get("desc") or it.get("item") or it.get("name") or ""
-        qty = it.get("quantity") or it.get("qty") or it.get("Quantity") or ""
-        unit = it.get("unit_price") or it.get("unitPrice") or it.get("price") or ""
-        amt = it.get("amount") or it.get("Line_total") or it.get("line_total") or it.get("total") or ""
-
-        # Extract item-level tax if available under common keys
-        tax_val = it.get("tax") or it.get("tax_amount") or it.get("line_tax") or it.get("item_tax") or it.get("taxAmount") or ""
-
-        # Extract explicit line total if present; otherwise fall back to amount
-        line_total_val = it.get("Line_total") or it.get("line_total") or it.get("lineTotal") or amt
-
-        item_rows.append({
-            "Description": str(desc).strip(),
-            "Quantity": float(clean_number(qty)),
-            "Unit Price": float(clean_number(unit)),
-            "Amount": float(clean_number(amt)),
-            "Tax": float(clean_number(tax_val)),
-            "Line Total": float(clean_number(line_total_val))
-        })
-
-    ui["Itemized Data"] = item_rows
+    # summary / totals
+    ui["Subtotal"] = clean_number(pick_first("subtotal", "sub_total", "sub total") or 0.0)
+    ui["Tax Percentage"] = clean_number(pick_first("tax_rate", "tax_percentage", "tax pct", "tax percentage") or 0.0)
+    ui["Total Tax"] = clean_number(pick_first("tax_amount", "tax", "total_tax") or 0.0)
+    ui["Total Amount"] = clean_number(pick_first("total_amount", "grand_total", "total", "amount") or 0.0)
+    ui["Currency"] = (pick_first("currency") or "").strip()
+
+    # Item extraction:
+    items_rows = []
+
+    # --- Primary approach: detect explicit list-of-dicts candidates first (unchanged) ---
+    def list_looks_like_items(lst):
+        if not isinstance(lst, list) or not lst:
+            return False
+        if not isinstance(lst[0], dict):
+            return False
+        # check if any expected item key present in first element
+        expected = {"descriptions", "description", "desc", "item", "quantity", "qty", "amount", "unit_price", "line_total", "line_total".lower(), "line_total"}
+        keys0 = {str(k).strip().lower() for k in lst[0].keys()}
+        return bool(expected.intersection(keys0))
+
+    for cand in list_candidates:
+        if list_looks_like_items(cand):
+            for it in cand:
+                row = map_item_dict(it)
+                if row is not None:
+                    items_rows.append(row)
+            # prefer first plausible list
+            if items_rows:
+                break
+
+    # --- Secondary approach: if parsed is a single dict that itself contains the item fields
+    # This is important because your model sometimes emits a single item as a top-level dict
+    # (e.g. {"descriptions":"...","quantity":"1.00","unit_price":"35,000.00",...}).
+    # We must map that directly (do NOT rely on finding a list named "items").
+    if not items_rows:
+        single_candidate_keys = {k.strip().lower() for k in (parsed.keys() if isinstance(parsed, dict) else [])} if isinstance(parsed, dict) else set()
+        # item-like keys we expect in the raw model output (explicitly include variants the model uses)
+        item_like_keys = {"descriptions", "description", "desc", "item", "quantity", "qty", "unit_price", "unit price", "price", "amount", "line_total", "line total", "line_total", "line_total".lower(), "sku", "tax", "tax_amount"}
+        if single_candidate_keys and single_candidate_keys.intersection(item_like_keys):
+            # map the parsed dict as a single line item
+            single_row = map_item_dict(parsed)
+            if single_row is not None:
+                items_rows.append(single_row)
+
+    # 2) If no list-of-dicts found, try to find a single dict anywhere that looks like an item (e.g., 'items': {...} as dict)
+    if not items_rows:
+        # search key_map values for dicts that have item-like keys
+        for k, vals in key_map.items():
+            for v in vals:
+                if isinstance(v, dict):
+                    # does this dict have an item-like key?
+                    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)
+                            # we don't break because there might be multiple item-like dicts at different keys,
+                            # but continue scanning to collect all.
+    # 3) Last resort: if key_map contains 'descriptions' or 'amount' as scalar but no dict, build a single-item row
+    if not items_rows:
+        desc = pick_first("descriptions", "description")
+        amt = pick_first("amount", "line_total")
+        qty = pick_first("quantity", "qty")
+        unit_price = pick_first("unit_price", "price")
+        if desc or amt or qty or unit_price:
+            items_rows.append({
+                "Description": str(desc or ""),
+                "Quantity": float(clean_number(qty)),
+                "Unit Price": float(clean_number(unit_price)),
+                "Amount": float(clean_number(amt)),
+                "Tax": float(clean_number(pick_first("tax", "tax_amount") or 0.0)),
+                "Line Total": float(clean_number(amt or 0.0))
+            })
+
+    ui["Itemized Data"] = items_rows
+
+    # Also set Sender/Recipient convenience fields
+    ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
+    ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
 
     return ui
 
@@ -425,13 +484,47 @@ def flatten_invoice_to_rows(invoice_data) -> list:
     """
     Converts nested invoice data into a flat list of rows (one per line item),
     with invoice-level and sender/recipient/bank fields repeated in each row.
+    This version collects bank details from both:
+      - invoice_data.get("Bank Details", {})  (nested dict style)
+      - top-level keys in invoice_data that start with 'bank_'
+    Ensures the expected bank_* columns always exist in the produced rows.
     """
+    EXPECTED_BANK_FIELDS = [
+        "bank_name",
+        "bank_account_number",
+        "bank_acc_name",
+        "bank_iban",
+        "bank_swift",
+        "bank_routing",
+        "bank_branch"
+    ]
+
     rows = []
-    line_items = invoice_data.get("Itemized Data", [])
+    invoice_data = invoice_data or {}
 
-    if not line_items:
-        # If no line items, create one row with invoice info only
-        row = {
+    # Collect line items (if present)
+    line_items = invoice_data.get("Itemized Data", []) or []
+
+    # Collect bank details from nested dict (if any) and from top-level bank_ keys
+    bank_details = {}
+    nested = invoice_data.get("Bank Details", {}) or {}
+    if isinstance(nested, dict):
+        for k, v in nested.items():
+            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
+            bank_details[key_name] = v
+
+    # also collect flat top-level bank_* keys (these come from your form_data)
+    for k, v in invoice_data.items():
+        if isinstance(k, str) and k.lower().startswith("bank_"):
+            bank_details[k] = v
+
+    # ensure all expected bank fields are present (empty string if missing)
+    for f in EXPECTED_BANK_FIELDS:
+        bank_details.setdefault(f, "")
+
+    # Helper to create base invoice row (shared for empty-items case and per-item rows)
+    def base_invoice_info():
+        return {
             "Invoice Number": invoice_data.get("Invoice Number", ""),
             "Invoice Date": invoice_data.get("Invoice Date", ""),
             "Due Date": invoice_data.get("Due Date", ""),
@@ -440,20 +533,19 @@ def flatten_invoice_to_rows(invoice_data) -> list:
             "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
             "Total Tax": invoice_data.get("Total Tax", 0.0),
             "Total Amount": invoice_data.get("Total Amount", 0.0),
-            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
-            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
-            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
-            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
+            "Sender Name": invoice_data.get("Sender Name", "") or (invoice_data.get("Sender",{}) or {}).get("Name",""),
+            "Sender Address": invoice_data.get("Sender Address", "") or (invoice_data.get("Sender",{}) or {}).get("Address",""),
+            "Recipient Name": invoice_data.get("Recipient Name", "") or (invoice_data.get("Recipient",{}) or {}).get("Name",""),
+            "Recipient Address": invoice_data.get("Recipient Address", "") or (invoice_data.get("Recipient",{}) or {}).get("Address",""),
         }
 
-        # Flatten bank details
-        bank = invoice_data.get("Bank Details", {})
-        for k, v in bank.items():
-            # Avoid double-prefixing if key already contains 'bank_'
-            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
-            row[key_name] = v
-
-        # Add empty line item fields
+    # If no line items, emit a single invoice-only row (with empty item columns)
+    if not line_items:
+        row = base_invoice_info()
+        # include all expected bank fields (consistent names)
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = bank_details.get(k, "")
+        # item columns (empty)
         row.update({
             "Item Description": "",
             "Item Quantity": 0,
@@ -465,40 +557,20 @@ def flatten_invoice_to_rows(invoice_data) -> list:
         rows.append(row)
         return rows
 
-    # For each line item, create a row with all invoice context
+    # For each line item, create a row with all invoice context + bank fields
     for item in line_items:
-        row = {
-            "Invoice Number": invoice_data.get("Invoice Number", ""),
-            "Invoice Date": invoice_data.get("Invoice Date", ""),
-            "Due Date": invoice_data.get("Due Date", ""),
-            "Currency": invoice_data.get("Currency", ""),
-            "Subtotal": invoice_data.get("Subtotal", 0.0),
-            "Tax Percentage": invoice_data.get("Tax Percentage", 0.0),
-            "Total Tax": invoice_data.get("Total Tax", 0.0),
-            "Total Amount": invoice_data.get("Total Amount", 0.0),
-            "Sender Name": invoice_data.get("Sender", {}).get("Name", ""),
-            "Sender Address": invoice_data.get("Sender", {}).get("Address", ""),
-            "Recipient Name": invoice_data.get("Recipient", {}).get("Name", ""),
-            "Recipient Address": invoice_data.get("Recipient", {}).get("Address", ""),
-        }
-
-        # Flatten bank details
-        bank = invoice_data.get("Bank Details", {})
-        for k, v in bank.items():
-            # Avoid double-prefixing if key already contains 'bank_'
-            key_name = k if str(k).startswith("bank_") else f"bank_{k}"
-            row[key_name] = v
-
-        # Add line item fields
+        row = base_invoice_info()
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = bank_details.get(k, "")
+        # try to read canonical item keys (safe .get)
         row.update({
-            "Item Description": item.get("Description", ""),
-            "Item Quantity": item.get("Quantity", 0),
-            "Item Unit Price": item.get("Unit Price", 0.0),
-            "Item Amount": item.get("Amount", 0.0),
-            "Item Tax": item.get("Tax", 0.0),
-            "Item Line Total": item.get("Line Total", item.get("Amount", 0.0)),
+            "Item Description": item.get("Description", "") if isinstance(item, dict) else "",
+            "Item Quantity": item.get("Quantity", 0) if isinstance(item, dict) else 0,
+            "Item Unit Price": item.get("Unit Price", 0.0) if isinstance(item, dict) else 0.0,
+            "Item Amount": item.get("Amount", 0.0) if isinstance(item, dict) else 0.0,
+            "Item Tax": item.get("Tax", 0.0) if isinstance(item, dict) else 0.0,
+            "Item Line Total": item.get("Line Total", item.get("Amount", 0.0)) if isinstance(item, dict) else 0.0,
         })
-
         rows.append(row)
 
     return rows
@@ -779,38 +851,121 @@ elif len(st.session_state.batch_results) > 0:
 
         # ---------- Sender / Recipient ----------
         with tabs[1]:
-            sender_info = form_data.get('Sender', {'Name': '', 'Address': ''})
-            recipient_info = form_data.get('Recipient', {'Name': '', 'Address': ''})
+            # 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_info['Name'] = st.text_input("Sender Name*", value=sender_info.get('Name', ''), key=f"sender_name_{selected_hash}")
-                sender_info['Address'] = st.text_area("Sender Address*", value=sender_info.get('Address', ''), key=f"sender_address_{selected_hash}")
-                recipient_info['Name'] = st.text_input("Recipient Name*", value=recipient_info.get('Name', ''), key=f"recipient_name_{selected_hash}")
-                recipient_info['Address'] = st.text_area("Recipient Address*", value=recipient_info.get('Address', ''), key=f"recipient_address_{selected_hash}")
+                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}")
+                
                 if st.button("⇄ Swap", help="Swap sender and recipient information", key=f"swap_{selected_hash}"):
-                    form_data['Sender'], form_data['Recipient'] = form_data['Recipient'], form_data['Sender']
+                    # 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']
                     st.rerun()
 
-        # ---------- Bank Details ----------
+        # ---------- Bank Details ---------- (FIXED)
+        # ---------- Bank Details ---------- (FIXED)
         with tabs[2]:
-            bank_info = form_data.get('Bank Details', {}) or {}
+            # Get bank details from nested dictionary
+            bank_details = form_data.get("Bank Details", {})
+            if not isinstance(bank_details, dict):
+                bank_details = {}
+            
+            bank_name = bank_details.get('bank_name', '')
+            bank_account_number = bank_details.get('bank_account_number', '')
+            bank_acc_name = bank_details.get('bank_acc_name', '')
+            bank_iban = bank_details.get('bank_iban', '')
+            bank_swift = bank_details.get('bank_swift', '')
+            bank_routing = bank_details.get('bank_routing', '')
+            bank_branch = bank_details.get('bank_branch', '')
+            
             with st.container():
-                bank_info['bank_name'] = st.text_input("Bank Name", value=bank_info.get('bank_name', ''), key=f"bank_name_{selected_hash}")
-                bank_info['bank_account_number'] = st.text_input("Account Number", value=bank_info.get('bank_account_number', '') or bank_info.get('bank_acc_no',''), key=f"bank_account_{selected_hash}")
-                bank_info['bank_acc_name'] = st.text_input("Bank Account Name", value=bank_info.get('bank_acc_name', '') or bank_info.get('bank_acc_name', ''), key=f"bank_acc_name_{selected_hash}")
-                bank_info['bank_iban'] = st.text_input("IBAN", value=bank_info.get('bank_iban', ''), key=f"iban_{selected_hash}")
-                bank_info['bank_swift'] = st.text_input("SWIFT Code", value=bank_info.get('bank_swift', ''), key=f"swift_code_{selected_hash}")
-                bank_info['bank_routing'] = st.text_input("Routing Number", value=bank_info.get('bank_routing', ''), key=f"routing_{selected_hash}")
-                bank_info['bank_branch'] = st.text_input("Branch", value=bank_info.get('bank_branch', ''), key=f"branch_{selected_hash}")
-            form_data['Bank Details'] = bank_info
+                bank_name = st.text_input("Bank Name", value=bank_name, key=f"bank_name_{selected_hash}")
+                bank_account_number = st.text_input("Account Number", value=bank_account_number, key=f"bank_account_{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_account_number"] = bank_account_number
+            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
 
         # ---------- Line Items ----------
+                # ---------- Line Items ----------
         with tabs[3]:
             editor_key = f"item_editor_{selected_hash}"
-            item_rows = form_data.get('Itemized Data', [])
-            df = pd.DataFrame(item_rows)
+            item_rows = form_data.get('Itemized Data', []) or []
+
+            # --- Normalize item keys produced by the model (e.g. "Item Description") into simple keys ---
+            def normalize_item_keys(item):
+                if not isinstance(item, dict):
+                    return {
+                        "Description": "",
+                        "Quantity": "",
+                        "Unit Price": "",
+                        "Amount": "",
+                        "Tax": "",
+                        "Line Total": ""
+                    }
+                mapping = {
+                    "Item Description": "Description",
+                    "description": "Description",
+                    "desc": "Description",
+                    "Item Quantity": "Quantity",
+                    "quantity": "Quantity",
+                    "qty": "Quantity",
+                    "Item Unit Price": "Unit Price",
+                    "unit_price": "Unit Price",
+                    "price": "Unit Price",
+                    "Item Amount": "Amount",
+                    "amount": "Amount",
+                    "Item Tax": "Tax",
+                    "tax": "Tax",
+                    "Item Line Total": "Line Total",
+                    "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] = ""
+
+                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] = ""
+
             st.write("✏️ Edit line items below. Press Enter or click outside a cell to confirm each edit.")
             edited_df = st.data_editor(
                 df,
@@ -821,6 +976,7 @@ 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)
         if st.button("💾 Save Edits for This File", key=f"save_{selected_hash}"):
             # Update line items from data editor