Create src/app.py

src/app.py

@@ -0,0 +1,1375 @@
+# =========================
+# Invoice Extractor (Qwen3-VL via RunPod vLLM) - Batch Mode with Tax Validation
+# =========================
+import os
+from pathlib import Path
+
+# -----------------------------
+# Environment hardening (HF Spaces, /.cache issue)
+# -----------------------------
+_home = os.environ.get("HOME", "")
+if _home in ("", "/", None):
+    repo_dir = os.getcwd()
+    safe_home = repo_dir if os.access(repo_dir, os.W_OK) else "/tmp"
+    os.environ["HOME"] = safe_home
+    print(f"[startup] HOME not set or unwritable — setting HOME={safe_home}")
+
+streamlit_dir = Path(os.environ["HOME"]) / ".streamlit"
+try:
+    streamlit_dir.mkdir(parents=True, exist_ok=True)
+    print(f"[startup] ensured {streamlit_dir}")
+except Exception as e:
+    print(f"[startup] WARNING: could not create {streamlit_dir}: {e}")
+
+# -----------------------------
+# Imports
+# -----------------------------
+import json
+from io import BytesIO
+import hashlib
+from typing import Dict, Any
+from datetime import datetime
+
+import streamlit as st
+import pandas as pd
+from PIL import Image
+
+# Optional: pdf2image is only needed for PDFs
+try:
+    from pdf2image import convert_from_bytes
+except Exception:
+    convert_from_bytes = None
+
+# -----------------------------
+# RunPod vLLM Configuration (from environment variables)
+# -----------------------------
+import requests
+import base64
+import re
+
+POD_URL = os.getenv("POD_URL", "")
+VLLM_API_KEY = os.getenv("VLLM_API_KEY", "")
+MODEL_NAME = "qwen3-vl-8b-nu-merged"
+
+# Validate secrets are set
+if not POD_URL or not VLLM_API_KEY:
+    st.error("⚠️ API credentials not configured. Please set POD_URL and VLLM_API_KEY in Space settings.")
+    st.stop()
+# -----------------------------
+# Page config & CSS
+# -----------------------------
+st.set_page_config(page_title="Invoice Extractor (Qwen3-VL) - Batch Mode", layout="wide")
+st.title("Invoice Extraction")
+
+st.markdown(
+    """
+    <style>
+        .stApp { background-color: #ECECEC !important; }
+        div.block-container { padding-top: 3rem; padding-bottom: 1rem; }
+        [data-testid="stSidebar"] { background-color: #F7F7F7 !important; }
+        div[data-testid="stTabs"] > div > div { padding-bottom: 6px !important; }
+        /* Keep right column steady on first render post-extraction */
+        [data-testid="column"]:nth-of-type(2) { min-height: 780px; }
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
+# Fixed sizes to prevent reflow wobble
+FIXED_IMG_WIDTH = 640
+DATA_EDITOR_HEIGHT = 380
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def ensure_state(k: str, default):
+    """Initialize a session_state key once, then let widgets bind to it via key=... (no value=...)."""
+    if k not in st.session_state:
+        st.session_state[k] = default
+
+def clean_float(x) -> float:
+    import re
+    if x is None:
+        return 0.0
+    if isinstance(x, (int, float)):
+        return float(x)
+    s = str(x).strip()
+    if s == "":
+        return 0.0
+    s = re.sub(r"[,\s]", "", s)
+    s = re.sub(r"[^\d\.\-]", "", s)
+    if s in ("", ".", "-", "-."):
+        return 0.0
+    try:
+        return float(s)
+    except Exception:
+        return 0.0
+
+def normalize_date(date_str) -> str:
+    """
+    Normalize various date formats to dd-MMM-yyyy format (e.g., 01-Jan-2025)
+    Handles: ISO, US, EU, and various other common date formats
+    Returns empty string if date cannot be parsed
+    """
+    if not date_str or date_str == "":
+        return ""
+
+    if isinstance(date_str, str):
+        date_str = date_str.strip()
+        if date_str == "":
+            return ""
+
+    # Common date formats to try
+    formats = [
+        "%Y-%m-%d",           # 2025-01-15 (ISO)
+        "%d-%m-%Y",           # 15-01-2025 (EU)
+        "%m-%d-%Y",           # 01-15-2025 (US)
+        "%Y/%m/%d",           # 2025/01/15
+        "%d/%m/%Y",           # 15/01/2025
+        "%m/%d/%Y",           # 01/15/2025
+        "%d.%m.%Y",           # 15.01.2025
+        "%Y.%m.%d",           # 2025.01.15
+        "%d %B %Y",           # 15 January 2025
+        "%d %b %Y",           # 15 Jan 2025
+        "%B %d, %Y",          # January 15, 2025
+        "%b %d, %Y",          # Jan 15, 2025
+        "%d-%b-%Y",           # 15-Jan-2025
+        "%d-%B-%Y",           # 15-January-2025
+        "%Y%m%d",             # 20250115
+    ]
+
+    parsed_date = None
+
+    # Try parsing with each format
+    for fmt in formats:
+        try:
+            parsed_date = datetime.strptime(str(date_str), fmt)
+            break
+        except (ValueError, TypeError):
+            continue
+
+    # If no format matched, return empty string
+    if parsed_date is None:
+        return ""
+
+    # Format as dd-MMM-yyyy (e.g., 01-Jan-2025)
+    return parsed_date.strftime("%d-%b-%Y")
+
+def parse_date_to_object(date_str):
+    """
+    Parse a date string to a datetime.date object for date_input widget
+    Returns None if date cannot be parsed
+    """
+    if not date_str or date_str == "":
+        return None
+
+    if isinstance(date_str, str):
+        date_str = date_str.strip()
+        if date_str == "":
+            return None
+
+    # Common date formats to try
+    formats = [
+        "%Y-%m-%d",           # 2025-01-15 (ISO)
+        "%d-%m-%Y",           # 15-01-2025 (EU)
+        "%m-%d-%Y",           # 01-15-2025 (US)
+        "%Y/%m/%d",           # 2025/01/15
+        "%d/%m/%Y",           # 15/01/2025
+        "%m/%d/%Y",           # 01/15/2025
+        "%d.%m.%Y",           # 15.01.2025
+        "%Y.%m.%d",           # 2025.01.15
+        "%d %B %Y",           # 15 January 2025
+        "%d %b %Y",           # 15 Jan 2025
+        "%B %d, %Y",          # January 15, 2025
+        "%b %d, %Y",          # Jan 15, 2025
+        "%d-%b-%Y",           # 15-Jan-2025
+        "%d-%B-%Y",           # 15-January-2025
+        "%Y%m%d",             # 20250115
+    ]
+
+    # Try parsing with each format
+    for fmt in formats:
+        try:
+            parsed_date = datetime.strptime(str(date_str), fmt)
+            return parsed_date.date()
+        except (ValueError, TypeError):
+            continue
+
+    return None
+
+# -----------------------------
+# HF login flow (REMOVED - No longer needed for vLLM API)
+# -----------------------------
+# Authentication is now handled via POD_URL and VLLM_API_KEY instead
+
+# -----------------------------
+# Model config
+# -----------------------------
+# OLD DONUT CODE (COMMENTED OUT - Now using vLLM API)
+# -----------------------------
+# HF_MODEL_ID = "Bhuvi13/model-V7"
+# TASK_PROMPT = "<s_cord-v2>"
+#
+# @st.cache_resource(show_spinner=False)
+# def load_model_and_processor(hf_model_id: str, task_prompt: str):
+#     ...
+
+# -----------------------------
+# vLLM Inference Function (RunPod API)
+# -----------------------------
+def run_inference_vllm(image: Image.Image):
+    """Run inference using RunPod vLLM API"""
+
+    # Extraction prompt (JSON format)
+    EXTRACTION_PROMPT = """Please carefully examine this invoice image and extract all the information into the following structured JSON format. Pay close attention to details and ensure accuracy in number formatting and text extraction.
+
+Extract the data into this exact JSON structure (do not add or remove keys):
+
+{
+  "header": {
+    "invoice_no": "Invoice number or reference ID",
+    "invoice_date": "Date the invoice was issued (maintain original format)",
+    "due_date": "Payment due date if specified",
+    "sender_name": "Name of the company/person issuing the invoice",
+    "sender_addr": "Complete address of the sender/issuer",
+    "rcpt_name": "Name of the recipient/customer",
+    "rcpt_addr": "Address of the recipient/customer",
+    "bank_iban": "International Bank Account Number",
+    "bank_name": "Name of the bank",
+    "bank_acc_no": "Bank account number",
+    "bank_routing": "Bank routing number",
+    "bank_swift": "SWIFT/BIC code",
+    "bank_acc_name": "Account holder name",
+    "bank_branch": "Bank branch information"
+  },
+  "items": [
+    {
+      "descriptions": "Detailed description of the item/service",
+      "SKU": "Stock Keeping Unit or item code",
+      "quantity": "Quantity of items",
+      "unit_price": "Price per unit",
+      "amount": "Total amount for this line item",
+      "tax": "Tax amount for this item",
+      "Line_total": "Total amount including tax for this line"
+    }
+  ],
+  "summary": {
+    "subtotal": "Subtotal amount before tax",
+    "tax_rate": "Tax rate percentage or description",
+    "tax_amount": "Total tax amount",
+    "total_amount": "Final total amount to be paid",
+    "currency": "Currency code (USD, EUR, etc.)"
+  }
+}
+
+STRICT POLICY RULES (apply exactly, do not deviate):
+1) Number formatting & types
+   - Preserve the original number formatting from the invoice (commas, decimal places, currency symbols in text fields if shown).
+   - In this JSON, output all values as strings. If a field is not present or cannot be determined with high confidence, output "" (empty string). Do not use null, 0, or placeholders.
+2) Currency selection (multi-currency invoices)
+   - If multiple currencies are shown, ALWAYS choose the recipient/customer currency for all monetary fields in items and summary.
+   - Do NOT perform FX conversion. Select the column/figures that are explicitly in the recipient’s currency.
+   - For "summary.currency", prefer the printed 3-letter code (e.g., USD, EUR, INR). If only an unambiguous symbol is present, map it (₹→INR, €→EUR, $→USD when clearly USD). If ambiguous, leave "".
+3) Tax handling (no rounding of rates; don’t recompute given totals)
+   - Do NOT round tax percentages. Use the original precision for any calculations; keep the printed formatting for "summary.tax_rate".
+   - If a TOTAL tax amount is explicitly printed on the invoice (e.g., “Tax”, “VAT”, “IGST”, “Total Tax”), TREAT IT AS AUTHORITATIVE. Do NOT recompute a new total.
+     a) If per-line tax amounts are printed, copy them directly.
+     b) If per-line tax amounts are not printed, allocate the printed TOTAL tax proportionally across line items by each line’s net amount (quantity * unit_price − discount). Use precise arithmetic; ensure the sum of allocated per-line taxes equals the printed TOTAL tax (adjust the last cent minimally if required).
+   - If NO total tax amount is printed but a tax rate is printed, compute per-line tax as: tax = (quantity * unit_price − discount) × (exact, unrounded tax rate). Then set "summary.tax_amount" = sum of per-line taxes.
+   - "items[].amount" is the pre-tax line amount AFTER discount. "items[].Line_total" = amount + tax.
+4) Discounts
+   - If discounts are present (per-line or overall), compute tax on the discounted base: (quantity * unit_price − discount). Never compute tax on the undiscounted amount.
+5) Due date calculation from payment terms
+   - Preserve the invoice’s original date format for both "invoice_date" and "due_date".
+   - If explicit due date is printed, use it as "due_date".
+   - If payment terms specify Net X (e.g., Net 30), set due_date = invoice_date + X days (same format as invoice_date).
+   - If terms say “upon receipt”, “upon publication”, or equivalent, due_date = invoice_date.
+   - If both a printed due date and terms exist and they conflict, prefer the printed due date.
+6) Items array
+   - Include every visible line item. Preserve multi-line descriptions using literal "\\n" where line breaks exist.
+   - If SKU is not shown, set "SKU": "".
+   - Ensure "quantity", "unit_price", "amount", "tax", and "Line_total" are consistent with the rules above.
+7) Summary invariants (when values are available on the invoice)
+   - "summary.subtotal" = sum of items[].amount.
+   - "summary.tax_amount" = sum of items[].tax (if you allocated or computed it). If the invoice prints a total tax amount, use that exact value and make per-line taxes sum to it.
+   - "summary.total_amount" = subtotal + tax_amount.
+   - If any of these values are not printed and cannot be derived reliably from the printed numbers, leave them as "".
+8) Text extraction fidelity
+   - Extract text exactly as printed (names, addresses, bank fields, references). Keep special characters and spacing (normalize only obvious OCR artifacts).
+   - If a bank field is absent (IBAN/SWIFT/routing/etc.), set it to "".
+
+Output constraints:
+- Return ONLY the JSON object described above (no explanations, no code fences, no trailing commas).
+- Keep all values as strings.
+- Do not add extra keys or sections beyond the given schema."""
+
+    try:
+        # Resize image if too large (max dimension 2048px to avoid payload size issues)
+        max_dimension = 2048
+        width, height = image.size
+        if width > max_dimension or height > max_dimension:
+            ratio = min(max_dimension / width, max_dimension / height)
+            new_size = (int(width * ratio), int(height * ratio))
+            image = image.resize(new_size, Image.Resampling.LANCZOS)
+            st.info(f"Image resized from {width}x{height} to {new_size[0]}x{new_size[1]} to reduce payload size")
+
+        # Convert image to base64
+        buffer = BytesIO()
+        image.save(buffer, format="PNG", optimize=True)
+        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+        # Check payload size
+        payload_size_mb = len(image_base64) / (1024 * 1024)
+        if payload_size_mb > 10:
+            st.warning(f"Warning: Large image payload ({payload_size_mb:.2f} MB). This might cause issues.")
+
+        data_url = f"data:image/png;base64,{image_base64}"
+
+        # Build payload
+        payload = {
+            "model": MODEL_NAME,
+            "messages": [
+                {"role": "system", "content": EXTRACTION_PROMPT},
+                {"role": "user", "content": [
+                    {"type": "image_url", "image_url": {"url": data_url}},
+                    {"type": "text", "text": "Extract invoice data."}
+                ]}
+            ],
+            "temperature": 0,
+            "max_tokens": 1536
+        }
+
+        headers = {
+            "Authorization": f"Bearer {VLLM_API_KEY}",
+            "Content-Type": "application/json"
+        }
+
+        # Call API
+        st.info(f"Sending request to API (payload size: {payload_size_mb:.2f} MB)...")
+        response = requests.post(
+            f"{POD_URL}/v1/chat/completions",
+            headers=headers,
+            json=payload,
+            timeout=90
+        )
+
+        if response.status_code == 200:
+            result = response.json()
+            return result["choices"][0]["message"]["content"]
+        else:
+            # Show detailed error for debugging
+            st.error(f"❌ API Error {response.status_code}")
+            try:
+                error_detail = response.json()
+                st.json(error_detail)  # Show as formatted JSON
+            except:
+                st.code(response.text)  # Show raw text
+            return None
+
+    except Exception as e:
+        st.error(f"Error calling vLLM: {str(e)}")
+        return None
+
+
+# -----------------------------
+# JSON Parser for vLLM Output
+# -----------------------------
+def parse_vllm_json(raw_json_text):
+    """Parse vLLM JSON output into structured format"""
+    try:
+        data = json.loads(raw_json_text)
+
+        def clean_amount(value):
+            if not value or value == "":
+                return 0.0
+            return float(re.sub(r'[^\d\.-]', '', str(value)))
+
+        header = data.get("header", {})
+        summary = data.get("summary", {})
+        items = data.get("items", [])
+
+        result = {
+            "Invoice Number": header.get("invoice_no", ""),
+            "Invoice Date": normalize_date(header.get("invoice_date", "")),
+            "Due Date": normalize_date(header.get("due_date", "")),
+            "Sender Name": header.get("sender_name", ""),
+            "Sender Address": header.get("sender_addr", ""),
+            "Sender": {
+                "Name": header.get("sender_name", ""),
+                "Address": header.get("sender_addr", "")
+            },
+            "Recipient Name": header.get("rcpt_name", ""),
+            "Recipient Address": header.get("rcpt_addr", ""),
+            "Recipient": {
+                "Name": header.get("rcpt_name", ""),
+                "Address": header.get("rcpt_addr", "")
+            },
+            "Bank Details": {
+                "bank_iban": header.get("bank_iban", ""),
+                "bank_name": header.get("bank_name", ""),
+                "bank_account_number": header.get("bank_acc_no", ""),
+                "bank_routing": header.get("bank_routing", ""),
+                "bank_swift": header.get("bank_swift", ""),
+                "bank_account_holder": header.get("bank_acc_name", ""),
+                "bank_branch": header.get("bank_branch", "")
+            },
+            "Subtotal": clean_amount(summary.get("subtotal", "0")),
+            "Tax Percentage": clean_amount(summary.get("tax_rate", "0")),
+            "Total Tax": clean_amount(summary.get("tax_amount", "0")),
+            "Total Amount": clean_amount(summary.get("total_amount", "0")),
+            "Currency": summary.get("currency", ""),
+            "Itemized Data": []
+        }
+
+        for item in items:
+            # Store raw tax value to distinguish empty ("") from explicit "0" or "0.00"
+            raw_tax = item.get("tax", "")
+
+            result["Itemized Data"].append({
+                "Description": item.get("descriptions", ""),
+                "SKU": item.get("SKU", ""),
+                "Quantity": clean_amount(item.get("quantity", "0")),
+                "Unit Price": clean_amount(item.get("unit_price", "0")),
+                "Amount": clean_amount(item.get("amount", "0")),
+                "Tax": clean_amount(raw_tax),
+                "Tax_Raw": raw_tax,  # Keep original to distinguish empty vs 0.00
+                "Line Total": clean_amount(item.get("Line_total", "0"))
+            })
+
+        return result
+
+    except Exception as e:
+        st.error(f"JSON parse error: {str(e)}")
+        return None
+
+
+# -----------------------------
+# Tax Validation Function
+# -----------------------------
+def validate_and_calculate_taxes(structured_data):
+    """
+    Enhanced tax validation with smart line-item calculation:
+    1. Calculate line-item tax ONLY when line item tax is empty/missing
+    2. Skip line items with explicit 0.00 tax (tax-exempt)
+    3. Skip validation if tax_amount is 0 but tax_rate exists
+    4. Ensure both Tax Percentage and Total Tax are properly filled
+    """
+
+    subtotal = structured_data.get("Subtotal", 0.0)
+    total_amount = structured_data.get("Total Amount", 0.0)
+    model_tax_rate = structured_data.get("Tax Percentage", 0.0)
+    model_tax_amount = structured_data.get("Total Tax", 0.0)
+    items = structured_data.get("Itemized Data", [])
+
+    # SKIP VALIDATION if: No tax detected (subtotal >= total) OR subtotal is invalid
+    if subtotal >= total_amount or subtotal <= 0:
+        structured_data["tax_validated"] = False
+        structured_data["tax_skip_reason"] = "No tax detected"
+        return structured_data
+
+    # SKIP if tax_rate exists but tax_amount is 0 (incomplete data)
+    if model_tax_rate > 0 and model_tax_amount == 0.0:
+        structured_data["tax_validated"] = False
+        structured_data["tax_skip_reason"] = "Tax rate exists but tax amount is 0"
+        return structured_data
+
+    # Determine authoritative tax rate from available sources
+    authoritative_rate = None
+    authority_source = None
+
+    # TEST SOURCE A: tax_rate
+    if model_tax_rate > 0:
+        expected_tax_from_rate = subtotal * (model_tax_rate / 100)
+        expected_total_from_rate = subtotal + expected_tax_from_rate
+        error_from_rate = abs(expected_total_from_rate - total_amount)
+    else:
+        error_from_rate = float('inf')
+
+    # TEST SOURCE B: tax_amount
+    if model_tax_amount > 0:
+        calculated_rate_from_amount = (model_tax_amount / subtotal) * 100
+        expected_total_from_amount = subtotal + model_tax_amount
+        error_from_amount = abs(expected_total_from_amount - total_amount)
+    else:
+        error_from_amount = float('inf')
+
+    # PICK WINNER (or use whichever is available)
+    if model_tax_rate > 0 or model_tax_amount > 0:
+        if error_from_rate < error_from_amount:
+            authoritative_rate = round(model_tax_rate, 4)
+            authority_source = "tax_rate"
+        else:
+            authoritative_rate = round(calculated_rate_from_amount, 4)
+            authority_source = "tax_amount"
+    else:
+        # No tax information available
+        structured_data["tax_validated"] = False
+        structured_data["tax_skip_reason"] = "No tax rate or amount provided"
+        return structured_data
+
+    # APPLY to line items - BUT respect explicit 0.00 values
+    calculated_total_tax = 0.0
+
+    for item in items:
+        amount = item.get("Amount", 0.0)
+        original_tax = item.get("Tax", 0.0)
+        raw_tax_value = item.get("Tax_Raw", "")  # Original string value from JSON
+
+        # If item amount is 0, keep tax at 0
+        if amount == 0.0:
+            item["Tax"] = 0.0
+            item["Line Total"] = 0.0
+            continue
+
+        # Distinguish between empty ("") and explicit "0" or "0.00"
+        # Empty string means tax was not provided - we should calculate it
+        # "0", "0.0", "0.00" means explicitly tax-exempt - skip calculation
+
+        is_explicitly_zero = False
+        if isinstance(raw_tax_value, str):
+            cleaned = raw_tax_value.strip()
+            # Check if it's explicitly set to some form of zero
+            if cleaned != "" and float(re.sub(r'[^\d\.-]', '', cleaned) or '0') == 0.0:
+                is_explicitly_zero = True
+        elif raw_tax_value == 0 or raw_tax_value == 0.0:
+            # If it's a number 0, treat as explicit
+            is_explicitly_zero = True
+
+        # If explicitly 0.00 - tax-exempt item, don't calculate
+        if is_explicitly_zero and original_tax == 0.0:
+            item["Tax"] = 0.0
+            item["Line Total"] = amount
+            calculated_total_tax += 0.0
+            continue
+
+        # Calculate tax for this line item
+        # Either: (1) Tax was empty/missing, or (2) Tax has a value that needs recalculation
+        corrected_tax = round(amount * (authoritative_rate / 100), 2)
+        item["Tax"] = corrected_tax
+        calculated_total_tax += corrected_tax
+        item["Line Total"] = round(amount + corrected_tax, 2)
+
+    # Update summary - ENSURE BOTH FIELDS ARE FILLED
+    structured_data["Tax Percentage"] = authoritative_rate
+    structured_data["Total Tax"] = round(calculated_total_tax, 2)
+    structured_data["Total Amount"] = round(subtotal + calculated_total_tax, 2)
+    structured_data["tax_validated"] = True
+    structured_data["tax_authority_source"] = authority_source
+    structured_data["original_tax_rate"] = model_tax_rate
+    structured_data["original_tax_amount"] = model_tax_amount
+
+    return structured_data
+
+
+# -----------------------------
+# ORIGINAL (previous) mapping logic — restored verbatim
+# -----------------------------
+def map_prediction_to_ui(pred):
+    import json, re
+    from collections import defaultdict
+
+    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:
+                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
+
+    def clean_number(x):
+        if x is None:
+            return 0.0
+        if isinstance(x, (int, float)):
+            return float(x)
+        s = str(x).strip()
+        if s == "":
+            return 0.0
+        s = re.sub(r"[,\s]", "", s)
+        s = re.sub(r"[^\d\.\-]", "", s)
+        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)
+        elif isinstance(obj, list):
+            for it in obj:
+                collect_keys(it, out)
+
+    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)
+
+    def map_item_dict(it):
+        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 ""
+        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))
+        }
+
+    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
+
+    ui = {
+        "Invoice Number": "",
+        "Invoice Date": "",
+        "Due Date": "",
+        "Currency": "",
+        "Subtotal": 0.0,
+        "Tax Percentage": 0.0,
+        "Total Tax": 0.0,
+        "Total Amount": 0.0,
+        "Sender": {"Name": "", "Address": ""},
+        "Recipient": {"Name": "", "Address": ""},
+        "Sender Name": "",
+        "Sender Address": "",
+        "Recipient Name": "",
+        "Recipient Address": "",
+        "Bank Details": {},
+        "Itemized Data": []
+    }
+
+    key_map = defaultdict(list)
+    list_candidates = []
+    if isinstance(parsed, dict):
+        collect_keys(parsed, key_map)
+        collect_lists_of_dicts(parsed, list_candidates)
+    elif isinstance(pred, dict):
+        collect_keys(pred, key_map)
+        collect_lists_of_dicts(pred, list_candidates)
+
+    def pick_first(*candidate_keys):
+        for k in candidate_keys:
+            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
+                    s = str(v).strip()
+                    if s != "":
+                        return s
+        return None
+
+    ui["Invoice Number"] = pick_first("invoice_no", "invoice_number", "invoiceid", "invoice id") or ""
+    ui["Invoice Date"] = normalize_date(pick_first("invoice_date", "date", "invoice date") or "")
+    ui["Due Date"] = normalize_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 = {}
+    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_", ""))
+        if val:
+            if bk == "iban":
+                bank["bank_iban"] = str(val)
+            else:
+                bank[bk if bk != "bank_acc_no" else "bank_account_number"] = str(val)
+    ui["Bank Details"] = bank
+
+    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()
+
+    items_rows = []
+
+    def list_looks_like_items(lst):
+        if not isinstance(lst, list) or not lst:
+            return False
+        if not isinstance(lst[0], dict):
+            return False
+        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)
+            if items_rows:
+                break
+
+    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 = {"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):
+            single_row = map_item_dict(parsed)
+            if single_row is not None:
+                items_rows.append(single_row)
+
+    if not items_rows:
+        for k, vals in key_map.items():
+            for v in vals:
+                if isinstance(v, dict):
+                    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 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
+    ui["Sender"] = {"Name": ui["Sender Name"], "Address": ui["Sender Address"]}
+    ui["Recipient"] = {"Name": ui["Recipient Name"], "Address": ui["Recipient Address"]}
+
+    return ui
+
+def flatten_invoice_to_rows(invoice_data) -> list:
+    EXPECTED_BANK_FIELDS = [
+        "bank_name",
+        "bank_account_number",
+        "bank_acc_name",
+        "bank_iban",
+        "bank_swift",
+        "bank_routing",
+        "bank_branch"
+    ]
+
+    # Helper to format text fields (empty -> NA)
+    def format_text_field(value):
+        if value is None or str(value).strip() == "":
+            return "NA"
+        return str(value).strip()
+
+    # Helper to format amount fields (empty -> 0)
+    def format_amount_field(value):
+        if value is None or value == "" or (isinstance(value, str) and value.strip() == ""):
+            return 0
+        try:
+            return float(value)
+        except (ValueError, TypeError):
+            return 0
+
+    rows = []
+    invoice_data = invoice_data or {}
+    line_items = invoice_data.get("Itemized Data", []) or []
+
+    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
+
+    for k, v in invoice_data.items():
+        if isinstance(k, str) and k.lower().startswith("bank_"):
+            bank_details[k] = v
+
+    for f in EXPECTED_BANK_FIELDS:
+        bank_details.setdefault(f, "")
+
+    def base_invoice_info():
+        return {
+            "Invoice Number": format_text_field(invoice_data.get("Invoice Number", "")),
+            "Invoice Date": format_text_field(invoice_data.get("Invoice Date", "")),
+            "Due Date": format_text_field(invoice_data.get("Due Date", "")),
+            "Currency": format_text_field(invoice_data.get("Currency", "")),
+            "Subtotal": format_amount_field(invoice_data.get("Subtotal", 0.0)),
+            "Tax Percentage": format_amount_field(invoice_data.get("Tax Percentage", 0.0)),
+            "Total Tax": format_amount_field(invoice_data.get("Total Tax", 0.0)),
+            "Total Amount": format_amount_field(invoice_data.get("Total Amount", 0.0)),
+            "Sender Name": format_text_field(invoice_data.get("Sender Name", "") or (invoice_data.get("Sender",{}) or {}).get("Name","")),
+            "Sender Address": format_text_field(invoice_data.get("Sender Address", "") or (invoice_data.get("Sender",{}) or {}).get("Address","")),
+            "Recipient Name": format_text_field(invoice_data.get("Recipient Name", "") or (invoice_data.get("Recipient",{}) or {}).get("Name","")),
+            "Recipient Address": format_text_field(invoice_data.get("Recipient Address", "") or (invoice_data.get("Recipient",{}) or {}).get("Address","")),
+        }
+
+    if not line_items:
+        row = base_invoice_info()
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = format_text_field(bank_details.get(k, ""))
+        row.update({
+            "Item Description": "NA",
+            "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
+
+    for item in line_items:
+        row = base_invoice_info()
+        for k in EXPECTED_BANK_FIELDS:
+            row[k] = format_text_field(bank_details.get(k, ""))
+        row.update({
+            "Item Description": format_text_field(item.get("Description", "") if isinstance(item, dict) else ""),
+            "Item Quantity": format_amount_field(item.get("Quantity", 0) if isinstance(item, dict) else 0),
+            "Item Unit Price": format_amount_field(item.get("Unit Price", 0.0) if isinstance(item, dict) else 0.0),
+            "Item Amount": format_amount_field(item.get("Amount", 0.0) if isinstance(item, dict) else 0.0),
+            "Item Tax": format_amount_field(item.get("Tax", 0.0) if isinstance(item, dict) else 0.0),
+            "Item Line Total": format_amount_field(item.get("Line Total", item.get("Amount", 0.0)) if isinstance(item, dict) else 0.0),
+        })
+        rows.append(row)
+    return rows
+
+# -----------------------------
+# Load model (COMMENTED OUT - Now using vLLM API)
+# -----------------------------
+# try:
+#     with st.spinner("Loading model & processor (cached) ..."):
+#         processor, model, device, decoder_input_ids = load_model_and_processor(HF_MODEL_ID, TASK_PROMPT)
+# except Exception as e:
+#     st.error("Could not load model automatically. See details below.")
+#     st.exception(e)
+#     st.stop()
+
+# -----------------------------
+# Session scaffolding
+# -----------------------------
+if "batch_results" not in st.session_state:
+    st.session_state.batch_results = {}
+if "current_file_hash" not in st.session_state:
+    st.session_state.current_file_hash = None
+if "is_processing_batch" not in st.session_state:
+    st.session_state.is_processing_batch = False
+
+# -----------------------------
+# Pre-mount two-column skeleton to avoid layout jump
+# -----------------------------
+frame_left, frame_right = st.columns([1, 1], vertical_alignment="top")
+
+# -----------------------------
+# Upload / Process
+# -----------------------------
+if not st.session_state.is_processing_batch and len(st.session_state.batch_results) == 0:
+    with frame_left:
+        st.header("📤 Upload Invoices")
+        uploaded_files = st.file_uploader(
+            "Upload invoice images (png/jpg/jpeg/pdf)",
+            type=["png", "jpg", "jpeg", "pdf"],
+            accept_multiple_files=True
+        )
+
+        if uploaded_files:
+            st.session_state.is_processing_batch = True
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+
+            for idx, uploaded_file in enumerate(uploaded_files):
+                status_text.text(f"Processing {idx+1}/{len(uploaded_files)}: {uploaded_file.name}")
+                uploaded_bytes = uploaded_file.read()
+                file_hash = hashlib.sha256(uploaded_bytes).hexdigest()
+
+                if file_hash in st.session_state.batch_results:
+                    progress_bar.progress((idx + 1) / len(uploaded_files))
+                    continue
+
+                # Load image (first page for PDFs)
+                image = None
+                is_pdf = uploaded_file.name.lower().endswith('.pdf') or (hasattr(uploaded_file, 'type') and uploaded_file.type == 'application/pdf')
+                if is_pdf:
+                    if convert_from_bytes is None:
+                        st.warning(f"PDF {uploaded_file.name} could not be rendered (pdf2image/poppler missing).")
+                        continue
+                    try:
+                        pages = convert_from_bytes(uploaded_bytes, dpi=200)
+                        if len(pages) > 0:
+                            image = pages[0].convert("RGB")
+                        else:
+                            st.warning(f"PDF {uploaded_file.name} has no pages.")
+                            continue
+                    except Exception:
+                        st.warning(f"Could not render PDF {uploaded_file.name}. Ensure 'pdf2image' and poppler are installed.")
+                        continue
+                else:
+                    try:
+                        image = Image.open(BytesIO(uploaded_bytes)).convert("RGB")
+                    except Exception:
+                        st.warning(f"Failed to open {uploaded_file.name}.")
+                        continue
+
+                if image is None:
+                    continue
+
+                # vLLM Inference + parsing + tax validation
+                try:
+                    # Call vLLM API
+                    raw_json = run_inference_vllm(image)
+
+                    if raw_json:
+                        # Parse JSON response
+                        parsed_data = parse_vllm_json(raw_json)
+
+                        if parsed_data:
+                            # Apply tax validation
+                            mapped = validate_and_calculate_taxes(parsed_data)
+                        else:
+                            st.warning(f"Failed to parse JSON for {uploaded_file.name}")
+                            mapped = {}
+                    else:
+                        st.warning(f"No response from vLLM for {uploaded_file.name}")
+                        mapped = {}
+
+                    pred = raw_json  # Store raw JSON for debugging
+                except Exception as e:
+                    st.warning(f"Error processing {uploaded_file.name}: {str(e)}")
+                    pred = None
+                    mapped = {}
+
+                safe_mapped = mapped if isinstance(mapped, dict) else {}
+
+                st.session_state.batch_results[file_hash] = {
+                    "file_name": uploaded_file.name,
+                    "image": image,
+                    "raw_pred": pred,
+                    "mapped_data": safe_mapped,
+                    "edited_data": safe_mapped.copy()
+                }
+
+                progress_bar.progress((idx + 1) / len(uploaded_files))
+
+            status_text.text("✅ All files processed!")
+            st.session_state.is_processing_batch = False
+            st.rerun()
+
+    with frame_right:
+        st.caption("Preview & editor will appear here after extraction.")
+
+elif len(st.session_state.batch_results) > 0:
+
+    # --------- Top row: All-results download + Back button ----------
+    with frame_left:
+        all_rows = []
+        for file_hash, result in st.session_state.batch_results.items():
+            rows = flatten_invoice_to_rows(result["edited_data"])
+            for r in rows:
+                r["Source File"] = result.get("file_name", file_hash)
+            all_rows.extend(rows)
+
+        if all_rows:
+            full_df = pd.DataFrame(all_rows)
+            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]
+            csv_bytes = full_df.to_csv(index=False).encode("utf-8")
+            st.download_button("📦 Download All Results (CSV)", csv_bytes,
+                               file_name="all_extracted_invoices.csv", mime="text/csv", key="download_all_csv")
+
+    with frame_right:
+        if st.button("⬅️ Back to Upload"):
+            st.session_state.batch_results.clear()
+            st.session_state.current_file_hash = None
+            st.session_state.is_processing_batch = False
+            st.rerun()
+
+    # --------- Selector ----------
+    with frame_left:
+        file_options = {f"{v['file_name']} ({k[:6]})": k for k, v in st.session_state.batch_results.items()}
+        selected_display = st.selectbox("Select invoice to view/edit:", options=list(file_options.keys()), index=0, key="file_selector")
+        selected_hash = file_options[selected_display]
+        if st.session_state.current_file_hash != selected_hash:
+            st.session_state.current_file_hash = selected_hash
+
+    current = st.session_state.batch_results[selected_hash]
+    image = current["image"]
+    form_data = current["edited_data"]
+
+    # --------- Initialize widget state - FORCE UPDATE from form_data ----------
+    bank = form_data.get("Bank Details", {}) if isinstance(form_data.get("Bank Details", {}), dict) else {}
+
+    # Always update state from form_data (don't use ensure_state which only sets if not exists)
+    st.session_state[f"Invoice Number_{selected_hash}"] = form_data.get('Invoice Number', '')
+
+    # Parse dates to date objects for date_input widgets
+    invoice_date_obj = parse_date_to_object(form_data.get('Invoice Date', ''))
+    due_date_obj = parse_date_to_object(form_data.get('Due Date', ''))
+    st.session_state[f"Invoice Date_{selected_hash}"] = invoice_date_obj
+    st.session_state[f"Due Date_{selected_hash}"] = due_date_obj
+
+    st.session_state[f"Currency_{selected_hash}"] = form_data.get('Currency', 'USD') or 'USD'
+    st.session_state[f"Currency_Custom_{selected_hash}"] = form_data.get('Currency', '') if form_data.get('Currency') not in ['USD','EUR','GBP','INR'] else ''
+    st.session_state[f"Subtotal_{selected_hash}"] = float(form_data.get('Subtotal', 0.0))
+    st.session_state[f"Tax Percentage_{selected_hash}"] = float(form_data.get('Tax Percentage', 0.0))
+    st.session_state[f"Total Tax_{selected_hash}"] = float(form_data.get('Total Tax', 0.0))
+    st.session_state[f"Total Amount_{selected_hash}"] = float(form_data.get('Total Amount', 0.0))
+    st.session_state[f"Sender Name_{selected_hash}"] = form_data.get('Sender Name', '')
+    st.session_state[f"Sender Address_{selected_hash}"] = form_data.get('Sender Address', '')
+    st.session_state[f"Recipient Name_{selected_hash}"] = form_data.get('Recipient Name', '')
+    st.session_state[f"Recipient Address_{selected_hash}"] = form_data.get('Recipient Address', '')
+    st.session_state[f"Bank_bank_name_{selected_hash}"] = bank.get('bank_name', '')
+    st.session_state[f"Bank_bank_account_number_{selected_hash}"] = bank.get('bank_account_number', '') or bank.get('bank_acc_no', '')
+    st.session_state[f"Bank_bank_acc_name_{selected_hash}"] = bank.get('bank_acc_name', '')
+    st.session_state[f"Bank_bank_iban_{selected_hash}"] = bank.get('bank_iban', '')
+    st.session_state[f"Bank_bank_swift_{selected_hash}"] = bank.get('bank_swift', '')
+    st.session_state[f"Bank_bank_routing_{selected_hash}"] = bank.get('bank_routing', '')
+    st.session_state[f"Bank_bank_branch_{selected_hash}"] = bank.get('bank_branch', '')
+
+    # --------- Display (no wobble) ----------
+    with frame_left:
+        st.image(image, caption=current["file_name"], width=FIXED_IMG_WIDTH)
+        st.write(f"**File Hash:** {selected_hash[:8]}...")
+        if current.get('raw_pred') is not None:
+            with st.expander("🔍 Show raw model output"):
+                st.json(current['raw_pred'])
+
+        if st.button("🔁 Re-Run Inference", key=f"rerun_{selected_hash}"):
+            with st.spinner("Re-running inference..."):
+                try:
+                    # Call vLLM API
+                    raw_json = run_inference_vllm(image)
+
+                    if raw_json:
+                        # Parse JSON response
+                        parsed_data = parse_vllm_json(raw_json)
+
+                        if parsed_data:
+                            # Apply tax validation
+                            mapped = validate_and_calculate_taxes(parsed_data)
+                        else:
+                            st.error("Failed to parse JSON response")
+                            mapped = {}
+                    else:
+                        st.error("No response from vLLM")
+                        mapped = {}
+
+                    safe_mapped = mapped if isinstance(mapped, dict) else {}
+                    pred = raw_json  # Store raw JSON
+
+                    # Update stored results
+                    st.session_state.batch_results[selected_hash]["raw_pred"] = pred
+                    st.session_state.batch_results[selected_hash]["mapped_data"] = mapped
+                    st.session_state.batch_results[selected_hash]["edited_data"] = safe_mapped.copy()
+
+                    # Clear widget state for this file so defaults refresh from new mapped data
+                    for key in [k for k in st.session_state.keys() if k.endswith(f"_{selected_hash}")]:
+                        del st.session_state[key]
+
+                    st.success("✅ Re-run complete")
+                    st.rerun()
+                except Exception as e:
+                    st.error(f"Re-run failed: {e}")
+
+    with frame_right:
+        st.subheader(f"Editable Invoice: {current['file_name']}")
+
+        # Quick swap outside the form (one clean rerun)
+        swap_cols = st.columns([1,1,2])
+        with swap_cols[0]:
+            if st.button("⇄ Swap Sender ↔ Recipient", key=f"swap_{selected_hash}"):
+                sn = f"Sender Name_{selected_hash}"
+                rn = f"Recipient Name_{selected_hash}"
+                sa = f"Sender Address_{selected_hash}"
+                ra = f"Recipient Address_{selected_hash}"
+                st.session_state[sn], st.session_state[rn] = st.session_state[rn], st.session_state[sn]
+                st.session_state[sa], st.session_state[ra] = st.session_state[ra], st.session_state[sa]
+                st.rerun()
+
+        # ----------------- FORM START -----------------
+        with st.form(key=f"edit_form_{selected_hash}", clear_on_submit=False):
+            tabs = st.tabs(["Invoice Details", "Sender/Recipient", "Bank Details", "Line Items"])
+
+            with tabs[0]:
+                st.text_input("Invoice Number", key=f"Invoice Number_{selected_hash}")
+                st.date_input("Invoice Date", key=f"Invoice Date_{selected_hash}", format="DD/MM/YYYY")
+                st.date_input("Due Date", key=f"Due Date_{selected_hash}", format="DD/MM/YYYY")
+
+                curr_options = ['USD', 'EUR', 'GBP', 'INR', 'Other']
+                if st.session_state[f"Currency_{selected_hash}"] not in curr_options:
+                    st.session_state[f"Currency_{selected_hash}"] = 'Other'
+                st.selectbox("Currency", options=curr_options, key=f"Currency_{selected_hash}")
+
+                if st.session_state.get(f"Currency_{selected_hash}") == 'Other':
+                    st.text_input("Specify Currency", key=f"Currency_Custom_{selected_hash}")
+
+                st.number_input("Subtotal", key=f"Subtotal_{selected_hash}")
+                st.number_input("Tax %", key=f"Tax Percentage_{selected_hash}")
+                st.number_input("Total Tax", key=f"Total Tax_{selected_hash}")
+                st.number_input("Total Amount", key=f"Total Amount_{selected_hash}")
+
+            with tabs[1]:
+                st.text_input("Sender Name", key=f"Sender Name_{selected_hash}")
+                st.text_area("Sender Address", key=f"Sender Address_{selected_hash}", height=80)
+                st.text_input("Recipient Name", key=f"Recipient Name_{selected_hash}")
+                st.text_area("Recipient Address", key=f"Recipient Address_{selected_hash}", height=80)
+
+            with tabs[2]:
+                st.text_input("Bank Name", key=f"Bank_bank_name_{selected_hash}")
+                st.text_input("Account Number", key=f"Bank_bank_account_number_{selected_hash}")
+                st.text_input("Account Name", key=f"Bank_bank_acc_name_{selected_hash}")
+                st.text_input("IBAN", key=f"Bank_bank_iban_{selected_hash}")
+                st.text_input("SWIFT", key=f"Bank_bank_swift_{selected_hash}")
+                st.text_input("Routing", key=f"Bank_bank_routing_{selected_hash}")
+                st.text_input("Branch", key=f"Bank_bank_branch_{selected_hash}")
+
+            with tabs[3]:
+                # Build base DF from current edited_data (not raw mapped) so it's always what the user last saved
+                item_rows = form_data.get('Itemized Data', []) or []
+                normalized = []
+                for it in item_rows:
+                    if not isinstance(it, dict):
+                        it = {}
+                    normalized.append({
+                        "Description": it.get("Description", it.get("Item Description", "")),
+                        "Quantity": it.get("Quantity", it.get("Item Quantity", 0)),
+                        "Unit Price": it.get("Unit Price", it.get("Item Unit Price", 0.0)),
+                        "Amount": it.get("Amount", it.get("Item Amount", 0.0)),
+                        "Tax": it.get("Tax", it.get("Item Tax", 0.0)),
+                        "Line Total": it.get("Line Total", it.get("Item Line Total", 0.0)),
+                    })
+
+                items_df = pd.DataFrame(normalized) if normalized else pd.DataFrame(
+                    columns=["Description", "Quantity", "Unit Price", "Amount", "Tax", "Line Total"]
+                )
+
+                # Show editor without totals
+                edited_df = st.data_editor(
+                    items_df,
+                    num_rows="dynamic",
+                    key=f"items_editor_{selected_hash}",
+                    use_container_width=True,
+                    height=DATA_EDITOR_HEIGHT - 50,  # Reduce height slightly for totals below
+                )
+
+                # Display non-editable totals row immediately below (looks integrated)
+                if len(edited_df) > 0:
+                    total_amount = edited_df["Amount"].sum()
+                    total_tax = edited_df["Tax"].sum()
+                    total_line_total = edited_df["Line Total"].sum()
+
+                    # Create totals display - styled to look like part of the table
+                    totals_df = pd.DataFrame([{
+                        "Description": "──── TOTAL ────",
+                        "Quantity": "",
+                        "Unit Price": "",
+                        "Amount": f"${total_amount:,.2f}",
+                        "Tax": f"${total_tax:,.2f}",
+                        "Line Total": f"${total_line_total:,.2f}"
+                    }])
+
+                    st.dataframe(
+                        totals_df,
+                        use_container_width=True,
+                        hide_index=True,
+                        height=38  # Single row height
+                    )
+
+            saved = st.form_submit_button("💾 Save All Edits")
+        # ----------------- FORM END -----------------
+
+        if saved:
+            currency = st.session_state.get(f"Currency_{selected_hash}", 'USD')
+            if currency == 'Other':
+                currency = st.session_state.get(f"Currency_Custom_{selected_hash}", '')
+
+            # Convert date objects to normalized strings (dd-MMM-yyyy format)
+            invoice_date = st.session_state.get(f"Invoice Date_{selected_hash}", None)
+            due_date = st.session_state.get(f"Due Date_{selected_hash}", None)
+
+            invoice_date_str = ""
+            if invoice_date is not None:
+                try:
+                    invoice_date_str = invoice_date.strftime("%d-%b-%Y")
+                except (AttributeError, ValueError):
+                    invoice_date_str = ""
+
+            due_date_str = ""
+            if due_date is not None:
+                try:
+                    due_date_str = due_date.strftime("%d-%b-%Y")
+                except (AttributeError, ValueError):
+                    due_date_str = ""
+
+            # Calculate totals from line items
+            line_items_list = edited_df.to_dict('records')
+            calculated_subtotal = sum(clean_float(item.get('Amount', 0)) for item in line_items_list)
+            calculated_total_tax = sum(clean_float(item.get('Tax', 0)) for item in line_items_list)
+            calculated_total = sum(clean_float(item.get('Line Total', 0)) for item in line_items_list)
+
+            # Calculate tax percentage if possible
+            calculated_tax_pct = 0.0
+            if calculated_subtotal > 0 and calculated_total_tax > 0:
+                calculated_tax_pct = round((calculated_total_tax / calculated_subtotal) * 100, 4)
+
+            updated = {
+                'Invoice Number': st.session_state.get(f"Invoice Number_{selected_hash}", ''),
+                'Invoice Date': invoice_date_str,
+                'Due Date': due_date_str,
+                'Currency': currency,
+                'Subtotal': calculated_subtotal,  # Auto-calculated from line items
+                'Tax Percentage': calculated_tax_pct,  # Auto-calculated
+                'Total Tax': calculated_total_tax,  # Auto-calculated from line items
+                'Total Amount': calculated_total,  # Auto-calculated from line items
+                '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_bank_name_{selected_hash}", ''),
+                    'bank_account_number': st.session_state.get(f"Bank_bank_account_number_{selected_hash}", ''),
+                    'bank_acc_name': st.session_state.get(f"Bank_bank_acc_name_{selected_hash}", ''),
+                    'bank_iban': st.session_state.get(f"Bank_bank_iban_{selected_hash}", ''),
+                    'bank_swift': st.session_state.get(f"Bank_bank_swift_{selected_hash}", ''),
+                    'bank_routing': st.session_state.get(f"Bank_bank_routing_{selected_hash}", ''),
+                    'bank_branch': st.session_state.get(f"Bank_bank_branch_{selected_hash}", '')
+                },
+                'Itemized Data': line_items_list,
+                'Sender': {"Name": st.session_state.get(f"Sender Name_{selected_hash}", ''),
+                           "Address": st.session_state.get(f"Sender Address_{selected_hash}", '')},
+                'Recipient': {"Name": st.session_state.get(f"Recipient Name_{selected_hash}", ''),
+                              "Address": st.session_state.get(f"Recipient Address_{selected_hash}", '')},
+            }
+
+            # Update session state fields to reflect the new calculated values
+            st.session_state[f"Subtotal_{selected_hash}"] = calculated_subtotal
+            st.session_state[f"Tax Percentage_{selected_hash}"] = calculated_tax_pct
+            st.session_state[f"Total Tax_{selected_hash}"] = calculated_total_tax
+            st.session_state[f"Total Amount_{selected_hash}"] = calculated_total
+
+            st.session_state.batch_results[selected_hash]["edited_data"] = updated
+            st.success(f"✅ Saved: {current['file_name']} | Updated totals: Subtotal=${calculated_subtotal:,.2f}, Tax=${calculated_total_tax:,.2f}, Total=${calculated_total:,.2f}")
+            st.rerun()  # Force rerun to show updated totals in the form fields
+
+        # Per-file CSV download (uses the current editor contents even if not saved)
+        d_currency = st.session_state.get(f"Currency_{selected_hash}", 'USD')
+        if d_currency == 'Other':
+            d_currency = st.session_state.get(f"Currency_Custom_{selected_hash}", '')
+
+        # Convert date objects to strings for download
+        d_invoice_date = st.session_state.get(f"Invoice Date_{selected_hash}", None)
+        d_due_date = st.session_state.get(f"Due Date_{selected_hash}", None)
+
+        d_invoice_date_str = ""
+        if d_invoice_date is not None:
+            try:
+                d_invoice_date_str = d_invoice_date.strftime("%d-%b-%Y")
+            except (AttributeError, ValueError):
+                d_invoice_date_str = ""
+
+        d_due_date_str = ""
+        if d_due_date is not None:
+            try:
+                d_due_date_str = d_due_date.strftime("%d-%b-%Y")
+            except (AttributeError, ValueError):
+                d_due_date_str = ""
+
+        download_data = {
+            'Invoice Number': st.session_state.get(f"Invoice Number_{selected_hash}", ''),
+            'Invoice Date': d_invoice_date_str,
+            'Due Date': d_due_date_str,
+            'Currency': d_currency,
+            'Subtotal': st.session_state.get(f"Subtotal_{selected_hash}", 0.0),
+            'Tax Percentage': st.session_state.get(f"Tax Percentage_{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_bank_name_{selected_hash}", ''),
+                'bank_account_number': st.session_state.get(f"Bank_bank_account_number_{selected_hash}", ''),
+                'bank_acc_name': st.session_state.get(f"Bank_bank_acc_name_{selected_hash}", ''),
+                'bank_iban': st.session_state.get(f"Bank_bank_iban_{selected_hash}", ''),
+                'bank_swift': st.session_state.get(f"Bank_bank_swift_{selected_hash}", ''),
+                'bank_routing': st.session_state.get(f"Bank_bank_routing_{selected_hash}", ''),
+                'bank_branch': st.session_state.get(f"Bank_bank_branch_{selected_hash}", '')
+            },
+            'Itemized Data': edited_df.to_dict('records')
+        }
+        rows = flatten_invoice_to_rows(download_data)
+        full_df = pd.DataFrame(rows)
+        csv_bytes_one = full_df.to_csv(index=False).encode("utf-8")
+        st.download_button(
+            "📥 Download This Invoice (CSV)",
+            csv_bytes_one,
+            file_name=f"{Path(current['file_name']).stem}_full.csv",
+            mime="text/csv",
+            key=f"dl_{selected_hash}"
+        )
+
+elif st.session_state.is_processing_batch:
+    with frame_left:
+        st.info("⏳ Processing batch... Please wait.")
+        st.progress(0)
+    with frame_right:
+        st.caption("Preview & editor will appear here after extraction.")
+
+else:
+    # Shouldn't happen, but keeps skeleton steady
+    with frame_left:
+        st.caption("Ready when you are.")
+    with frame_right:
+        st.caption("Preview & editor will appear here after extraction.")