Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import io | |
| import time | |
| import gc | |
| import numpy as np | |
| import pandas as pd | |
| from flask import Flask, request, send_file, jsonify | |
| from flask_cors import CORS | |
| import google.generativeai as genai | |
| from groq import Groq | |
| from PyPDF2 import PdfReader | |
| from dotenv import load_dotenv | |
| # OCR Imports (EasyOCR) | |
| import easyocr | |
| from pdf2image import convert_from_bytes, pdfinfo_from_bytes | |
| app = Flask(__name__) | |
| CORS(app, resources={r"/*": {"origins": "*"}}) | |
| # Load .env locally | |
| basedir = os.path.abspath(os.path.dirname(__file__)) | |
| load_dotenv(os.path.join(basedir, ".env")) | |
| # API Keys | |
| GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY") | |
| GROQ_API_KEY = os.environ.get("GROQ_API_KEY") | |
| print(f"DEBUG - Is Groq Key loaded? {'YES' if GROQ_API_KEY else 'NO'}") | |
| genai.configure(api_key=GEMINI_API_KEY) | |
| groq_client = Groq(api_key=GROQ_API_KEY) | |
| # EasyOCR reader (load once) | |
| easyocr_reader = easyocr.Reader(["en"], gpu=False) | |
| # OCR settings | |
| OCR_DPI = 200 | |
| # Groq batching settings | |
| BATCH_SIZE = 3 | |
| MAX_RETRIES = 5 | |
| SYSTEM_PROMPT = """ | |
| You are a data extraction assistant. Extract the payroll data from the provided document or text. | |
| Return ONLY a valid JSON object with a single key called "data". The value of "data" must be an array of objects. | |
| Do not include markdown formatting, code blocks, or explanations. | |
| Each object in the array must have exactly these keys: "SSN", "Participant", "Amount", "Contract", "Carrier Name". | |
| If a specific field is missing for a participant, use null. | |
| Example format: | |
| { | |
| "data": [ | |
| {"SSN": "123456789", "Participant": "John Doe", "Amount": "150.00", "Contract": null, "Carrier Name": "MG Trust Co"} | |
| ] | |
| } | |
| """ | |
| def home(): | |
| return jsonify({"status": "API running"}), 200 | |
| def generate_excel_buffer(json_data): | |
| """Convert JSON extracted data into an Excel file buffer.""" | |
| records = json_data.get("data", []) | |
| df = pd.DataFrame(records) | |
| expected_columns = ["SSN", "Participant", "Amount", "Contract", "Carrier Name"] | |
| for col in expected_columns: | |
| if col not in df.columns: | |
| df[col] = None | |
| df = df[expected_columns] | |
| output = io.BytesIO() | |
| with pd.ExcelWriter(output, engine="openpyxl") as writer: | |
| df.to_excel(writer, index=False, sheet_name="Payroll Data") | |
| output.seek(0) | |
| return output | |
| def extract_text_with_pypdf2(pdf_file): | |
| """Extract text from PDF using PyPDF2.""" | |
| pdf_reader = PdfReader(pdf_file) | |
| extracted_text = "" | |
| for page in pdf_reader.pages: | |
| text = page.extract_text() | |
| if text: | |
| extracted_text += text + "\n" | |
| return extracted_text.strip() | |
| def extract_text_with_easyocr_pagewise(pdf_bytes): | |
| """Extract OCR text page by page using EasyOCR (memory safe).""" | |
| info = pdfinfo_from_bytes(pdf_bytes) | |
| total_pages = info["Pages"] | |
| pages_text = [] | |
| for i in range(1, total_pages + 1): | |
| print(f"🔍 OCR Processing Page {i}/{total_pages}") | |
| images = convert_from_bytes( | |
| pdf_bytes, | |
| dpi=OCR_DPI, | |
| first_page=i, | |
| last_page=i | |
| ) | |
| img = images[0] | |
| img_np = np.array(img) | |
| result = easyocr_reader.readtext(img_np, detail=0, paragraph=True) | |
| page_text = "\n".join(result).strip() | |
| pages_text.append(page_text) | |
| # Free memory | |
| del images | |
| del img | |
| del img_np | |
| gc.collect() | |
| return pages_text | |
| def call_groq_with_retry(text_block): | |
| """Call Groq with retry support (handles 429).""" | |
| for attempt in range(MAX_RETRIES): | |
| try: | |
| chat_completion = groq_client.chat.completions.create( | |
| messages=[ | |
| {"role": "system", "content": SYSTEM_PROMPT.replace("\xa0", " ")}, | |
| {"role": "user", "content": f"Extract payroll rows ONLY from this text:\n\n{text_block}"}, | |
| ], | |
| model="llama-3.1-8b-instant", | |
| response_format={"type": "json_object"}, | |
| ) | |
| response_text = chat_completion.choices[0].message.content | |
| clean_text = response_text.replace("\xa0", " ").strip() | |
| return json.loads(clean_text) | |
| except Exception as e: | |
| error_msg = str(e) | |
| if "429" in error_msg: | |
| wait_time = (attempt + 1) * 3 | |
| print(f"⚠️ Groq Rate Limit Hit (429). Waiting {wait_time}s...") | |
| time.sleep(wait_time) | |
| continue | |
| raise e | |
| raise Exception("Groq API failed after multiple retries due to rate limit.") | |
| def clean_rows(rows): | |
| """ | |
| Merge duplicate rows smartly. | |
| If same Participant+Amount appears again and SSN is missing in old row, | |
| update it using new row. | |
| """ | |
| merged = {} | |
| for row in rows: | |
| ssn = str(row.get("SSN", "")).strip() if row.get("SSN") else "" | |
| participant = str(row.get("Participant", "")).strip() if row.get("Participant") else "" | |
| amount = str(row.get("Amount", "")).strip() if row.get("Amount") else "" | |
| contract = str(row.get("Contract", "")).strip() if row.get("Contract") else "" | |
| carrier = str(row.get("Carrier Name", "")).strip() if row.get("Carrier Name") else "" | |
| # skip completely empty rows | |
| if not participant and not ssn: | |
| continue | |
| # Main key: participant + amount (because SSN may appear later) | |
| key = (participant.lower(), amount) | |
| if key not in merged: | |
| merged[key] = { | |
| "SSN": ssn if ssn else None, | |
| "Participant": participant if participant else None, | |
| "Amount": amount if amount else None, | |
| "Contract": contract if contract else None, | |
| "Carrier Name": carrier if carrier else None, | |
| } | |
| else: | |
| # merge missing fields | |
| if not merged[key]["SSN"] and ssn: | |
| merged[key]["SSN"] = ssn | |
| if not merged[key]["Contract"] and contract: | |
| merged[key]["Contract"] = contract | |
| if not merged[key]["Carrier Name"] and carrier: | |
| merged[key]["Carrier Name"] = carrier | |
| # if participant missing (rare case) | |
| if not merged[key]["Participant"] and participant: | |
| merged[key]["Participant"] = participant | |
| return list(merged.values()) | |
| def process_with_gemini(): | |
| """Endpoint for Gemini (Image Model) PDF processing.""" | |
| if "file" not in request.files: | |
| return jsonify({"error": "No file uploaded"}), 400 | |
| file = request.files["file"] | |
| if file.filename == "": | |
| return jsonify({"error": "No selected file"}), 400 | |
| try: | |
| file_bytes = file.read() | |
| mime_type = file.mimetype | |
| model = genai.GenerativeModel("gemini-2.5-flash") | |
| contents = [ | |
| SYSTEM_PROMPT, | |
| {"mime_type": mime_type, "data": file_bytes}, | |
| ] | |
| response = model.generate_content(contents) | |
| response_text = ( | |
| response.text.strip() | |
| .removeprefix("```json") | |
| .removesuffix("```") | |
| .strip() | |
| ) | |
| parsed_data = json.loads(response_text) | |
| excel_buffer = generate_excel_buffer(parsed_data) | |
| return send_file( | |
| excel_buffer, | |
| mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", | |
| as_attachment=True, | |
| download_name="Extracted_Data_Gemini.xlsx", | |
| ) | |
| except Exception as e: | |
| return jsonify({"error": str(e)}), 500 | |
| def process_with_groq(): | |
| """Groq endpoint with OCR fallback + batching to avoid 429.""" | |
| if "file" not in request.files: | |
| return jsonify({"error": "No file uploaded"}), 400 | |
| file = request.files["file"] | |
| if file.filename == "": | |
| return jsonify({"error": "No selected file"}), 400 | |
| try: | |
| pdf_bytes = file.read() | |
| # Reset pointer for PyPDF2 | |
| file.seek(0) | |
| extracted_text = extract_text_with_pypdf2(file) | |
| all_rows = [] | |
| # Case 1: Text PDF | |
| if len(extracted_text.strip()) > 150: | |
| print("✅ Text detected using PyPDF2. Sending full document to Groq...") | |
| parsed_data = call_groq_with_retry(extracted_text) | |
| all_rows.extend(parsed_data.get("data", [])) | |
| # Case 2: Scanned PDF | |
| else: | |
| print("⚠️ No proper text found. Running EasyOCR page-by-page...") | |
| pages_text = extract_text_with_easyocr_pagewise(pdf_bytes) | |
| # Batch pages to reduce Groq calls | |
| for i in range(0, len(pages_text), BATCH_SIZE): | |
| batch_pages = pages_text[i:i + BATCH_SIZE] | |
| batch_text = "\n\n--- PAGE BREAK ---\n\n".join(batch_pages).strip() | |
| if len(batch_text) < 50: | |
| continue | |
| print(f"📤 Sending Pages {i+1} to {min(i+BATCH_SIZE, len(pages_text))} to Groq...") | |
| parsed_batch = call_groq_with_retry(batch_text) | |
| all_rows.extend(parsed_batch.get("data", [])) | |
| all_rows = clean_rows(all_rows) | |
| if not all_rows: | |
| return jsonify({"error": "No payroll rows extracted from document."}), 400 | |
| final_json = {"data": all_rows} | |
| excel_buffer = generate_excel_buffer(final_json) | |
| return send_file( | |
| excel_buffer, | |
| mimetype="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", | |
| as_attachment=True, | |
| download_name="Extracted_Data_Groq.xlsx", | |
| ) | |
| except Exception as e: | |
| print(f"\n--- CRITICAL BACKEND ERROR --- \n{str(e)}\n------------------------------\n") | |
| return jsonify({"error": str(e)}), 500 | |
| if __name__ == "__main__": | |
| port = int(os.environ.get("PORT", 7860)) | |
| app.run(host="0.0.0.0", port=port) |