Update app.py

app.py

@@ -1,132 +1,157 @@
 import os
+import re
+import json
 import gradio as gr
+import pandas as pd
 import pdfplumber
 import pytesseract
-from PIL import Image
 from pdf2image import convert_from_path
-import pandas as pd
-import numpy as np
-import re
+from huggingface_hub import InferenceClient
+
+# Initialize Hugging Face Inference Client
+hf_token = os.getenv("HF_TOKEN")
+client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3", token=hf_token)
 
-# For Excel files
 def extract_excel_data(file_path):
+    """Extract text from Excel file"""
     df = pd.read_excel(file_path, engine='openpyxl')
-    return df.to_string()
+    return df.to_string(index=False)
 
-# For PDF files with fallback OCR
 def extract_text_from_pdf(pdf_path, is_scanned=False):
+    """Extract text from PDF with fallback OCR"""
     try:
-        # First try native PDF extraction
+        # Try native PDF extraction first
         with pdfplumber.open(pdf_path) as pdf:
             text = ""
             for page in pdf.pages:
                 text += page.extract_text() + "\n"
             return text
     except Exception as e:
-        # Fallback to OCR if PDF is invalid
         print(f"Native PDF extraction failed: {str(e)}")
-        print("Trying OCR fallback...")
+        # Fallback to OCR for scanned PDFs
         images = convert_from_path(pdf_path, dpi=200)
         text = ""
         for image in images:
             text += pytesseract.image_to_string(image) + "\n"
         return text
 
-# Prompt engineering for structured extraction
 def parse_bank_statement(text):
-    # Clean up text from PDF/OCR artifacts
+    """Parse bank statement using LLM with fallback to rule-based parser"""
     cleaned_text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
     
+    # Craft precise prompt for LLM
     prompt = f"""
-    Extract the following details from the bank statement text:
-    - Transaction Date
-    - Description / Merchant
-    - Amount
-    - Debit / Credit
-    - Closing Balance
-    - Expense Type (if available)
+You are a financial data parser. Extract transactions from bank statements.
 
-    Return the results in JSON format with keys: 
-    ["date", "description", "amount", "debit_credit", "closing_balance", "expense_type"].
+Given this bank statement text:
 
-    Example:
-    {{
-      "transactions": [
-        {{
-          "date": "2025-06-01",
-          "description": "Grocery Store",
-          "amount": "150.00",
-          "debit_credit": "Debit",
-          "closing_balance": "1200.00",
-          "expense_type": "Food"
-        }}
-      ]
-    }}
+Extract all transactions with these fields:
+- Date
+- Description
+- Amount
+- Debit
+- Credit
+- Closing Balance
+- Category
 
-    Bank Statement Text:
-    {cleaned_text}
-    """
-    
-    # Simulate LLM response with deterministic parsing for demo
-    # Replace this with actual LLM inference in production
-    return simulate_llm_parsing(cleaned_text)
+Return JSON with "transactions" array containing these fields.
 
-def simulate_llm_parsing(text):
-    """Mock LLM response for demo purposes"""
-    # Simple regex-based parsing for demonstration
-    transactions = []
-    lines = text.split('\n')
-    
-    # Skip header lines
-    data_lines = lines[lines.index('Date') + 1:]
+Example format:
+{"transactions": [
+  {"date": "2025-05-08", "description": "Company XYZ Payroll", "amount": "8315.40", "debit": "0.00", "credit": "8315.40", "closing_balance": "38315.40", "category": "Salary"},
+  ...
+]}
+
+Rules:
+1. Ensure numeric fields have valid numbers (e.g., "0.00" instead of "-")
+2. Convert negative balances to standard format (e.g., "-2421.72")
+3. Map category names consistently (e.g., "Groceries", "Medical", "Utilities")
+"""
+
+    try:
+        # Call LLM via Hugging Face Inference API
+        response = client.text_generation(prompt, max_new_tokens=1000, temperature=0.1)
+        return json.loads(response)
+    except Exception as e:
+        print(f"LLM Error: {str(e)}")
+        # Fallback to rule-based parser
+        return rule_based_parser(cleaned_text)
+
+def rule_based_parser(text):
+    """Fallback parser for structured tables with pipe delimiters"""
+    lines = [line.strip() for line in text.split('\n') if line.strip()]
     
-    for i in range(0, len(data_lines), 7):  # Process in chunks of 7
-        if i+6 >= len(data_lines):
+    # Find header line containing 'Date'
+    header_index = None
+    for i, line in enumerate(lines):
+        if re.search(r'\bDate\b', line):
+            header_index = i
             break
+    
+    if header_index is None or header_index + 1 >= len(lines):
+        return {"transactions": []}
+    
+    data_lines = lines[header_index + 1:]
+    transactions = []
+    
+    for line in data_lines:
+        if not line.startswith('|'):
+            continue
+            
+        parts = [p.strip() for p in line.split('|') if p.strip()]
+        if len(parts) < 7:
+            continue
             
         try:
             transactions.append({
-                "date": data_lines[i].strip(),
-                "description": data_lines[i+1].strip(),
-                "amount": data_lines[i+2].strip(),
-                "debit_credit": data_lines[i+3].strip(),
-                "closing_balance": data_lines[i+5].strip(),
-                "expense_type": data_lines[i+6].strip()
+                "date": parts[0],
+                "description": parts[1],
+                "amount": parts[2],
+                "debit": parts[3],
+                "credit": parts[4],
+                "closing_balance": parts[5],
+                "category": parts[6]
             })
         except Exception as e:
-            print(f"Error parsing line {i}: {str(e)}")
-            continue
-            
+            print(f"Error parsing line: {str(e)}")
+    
     return {"transactions": transactions}
 
-# Main function
 def process_file(file, is_scanned):
+    """Main processing function"""
+    if not file:
+        return "No file uploaded"
+    
     file_path = file.name
     file_ext = os.path.splitext(file_path)[1].lower()
     
-    if file_ext == '.xlsx':
-        text = extract_excel_data(file_path)
-    elif file_ext == '.pdf':
-        text = extract_text_from_pdf(file_path, is_scanned=is_scanned)
-    else:
-        return "Unsupported file format. Please upload PDF or Excel."
-    
-    parsed_data = parse_bank_statement(text)
-    
-    # Convert to DataFrame for display
-    df = pd.DataFrame(parsed_data["transactions"])
-    return df
+    try:
+        if file_ext == '.xlsx':
+            text = extract_excel_data(file_path)
+        elif file_ext == '.pdf':
+            text = extract_text_from_pdf(file_path, is_scanned=is_scanned)
+        else:
+            return {"error": "Unsupported file format"}
+        
+        parsed_data = parse_bank_statement(text)
+        df = pd.DataFrame(parsed_data["transactions"])
+        return df
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Gradio interface
+# Gradio Interface
 interface = gr.Interface(
     fn=process_file,
     inputs=[
-        gr.File(label="Upload PDF/Excel"),
-        gr.Checkbox(label="Is Scanned PDF?")
+        gr.File(label="Upload Bank Statement (PDF/Excel)"),
+        gr.Checkbox(label="Is Scanned PDF? (Use OCR)")
     ],
-    outputs=gr.Dataframe(label="Extracted Transactions"),
-    title="Bank Statement Parser",
-    description="Convert PDF/Excel bank statements into structured data using hybrid parsing techniques.",
+    outputs=gr.Dataframe(
+        label="Parsed Transactions",
+        headers=["Date", "Description", "Amount", "Debit", "Credit", "Closing Balance", "Category"]
+    ),
+    title="AI Bank Statement Parser",
+    description="Extract structured transaction data from PDF/Excel bank statements using LLM and hybrid parsing techniques.",
     allow_flagging="never"
 )