"""Invoice extraction model worker on HF ZeroGPU. Public Space, but it exposes only a status panel. All invoice data lives in the owner's Cloudflare D1/R2; this worker claims jobs from the invoice-api Worker's internal API (secret-authenticated), extracts with LFM2-1.2B-Extract on a ZeroGPU slice, and posts results back. Space secrets: WORKER_URL, INTERNAL_KEY. """ import json import logging import os import re import threading import time import fitz # PyMuPDF import gradio as gr import httpx import spaces import torch from transformers import AutoModelForCausalLM, AutoTokenizer logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") log = logging.getLogger("worker") MODEL_ID = "LiquidAI/LFM2-1.2B-Extract" POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "3")) MAX_NEW_TOKENS = 2048 _TEMPLATE = """{ "invoice_number": null, "invoice_date": null, "due_date": null, "vendor": {"name": null, "address": null, "gstin": null, "pan": null, "email": null, "phone": null}, "buyer": {"name": null, "address": null, "gstin": null, "pan": null}, "service_description": null, "billing_period": null, "sac_code": null, "work_order": null, "amounts": {"taxable_value": null, "cgst": null, "sgst": null, "igst": null, "grand_total": null}, "amount_in_words": null, "bank_details": {"bank_name": null, "account_number": null, "ifsc": null, "micr": null}, "employees": [ {"name": null, "monthly_billing": null, "payable_days": null, "amount": null, "gst_amount": null, "total": null} ], "other_details": null }""" def build_system_prompt() -> str: if VENDORS: vendor_context = "\n".join( f"{i}. {v.get('name')} - GSTIN {v.get('gstin')}, PAN {v.get('pan')}. Bill numbers start with {v.get('prefix')}/." for i, v in enumerate(VENDORS, 1) ) else: vendor_context = os.environ.get("VENDOR_CONTEXT", "").strip() vc = f"\n\nKnown vendors (the invoice is always ISSUED BY exactly one of these):\n{vendor_context}\n" if vendor_context else "" return f"""You are an expert data extraction engine for Indian GST invoices. Fill this exact JSON template from the invoice text. Respond with ONLY the completed JSON object. {_TEMPLATE} {vc} Field guide: - "invoice_number": the value printed after "Bill No.:" (e.g. "XXX/003/26-27"). Never an employee code. - "vendor": the party that ISSUED the invoice (its name is in the letterhead at the very top). The GSTIN/PAN printed on the RIGHT side near "Bill No." belong to the VENDOR. - "buyer": the party being billed - the name/address block at the top LEFT. The GSTIN/PAN printed directly under the buyer's address belong to the BUYER. - "amounts": rupee amounts as plain JSON numbers, no commas, no symbols. Taxes (CGST/SGST/IGST) are amounts, not percentages; use null for any tax type not charged. "grand_total" is the final payable total (equals the amount in words). - "amount_in_words": the sentence after "Rupees:" exactly as printed. - "bank_details": the vendor's bank account printed under "Bank Details". - "employees": one entry per employee row in the annexure table at the end (ignore the TOTAL row). IMPORTANT: the first column is the manager ("Kind Attention Person") - the employee's own name is in the "Employee Name" column. "monthly_billing" = the Billing column (monthly rate), "payable_days" = Total Payable Days, "amount" = Total Payable Billing, "gst_amount" = the row's GST/IGST amount if shown, "total" = the last number in the row (row grand total). Use [] if there is no employee table. - "other_details": PO/work-order references, declarations, PF/ESIC numbers, or anything else notable. - Copy values exactly as printed; use null when absent. NEVER invent or guess values. - GSTIN is 15 characters; PAN is 10 characters.""" state = { "started_at": int(time.time()), "last_poll": None, "processed": 0, "failed": 0, "current_job": None, "last_error": None, } from anchors import VENDORS, deterministic_fields, merge_result # noqa: E402 class NoTextLayerError(Exception): pass # ---------- model ---------- log.info("Loading %s ...", MODEL_ID) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16) model.to("cuda") # ZeroGPU: safe at startup, GPU attaches inside @spaces.GPU calls log.info("Model ready.") @spaces.GPU(duration=120) def llm_generate(text: str, company_hint: str | None) -> str: hint = f"(Hint: this invoice relates to {company_hint}.)\n\n" if company_hint else "" messages = [ {"role": "system", "content": build_system_prompt()}, {"role": "user", "content": f"{hint}Invoice text:\n\n{text}"}, ] inputs = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) with torch.inference_mode(): out = model.generate( inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, pad_token_id=tokenizer.eos_token_id, ) return tokenizer.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True) _NUMERIC_KEYS = { "taxable_value", "cgst", "sgst", "igst", "grand_total", "monthly_billing", "payable_days", "amount", "gst_amount", "total", "quantity", "rate", } def _coerce_numbers(obj, key=None): """Recursively convert '1,26,166' style strings to numbers on numeric fields.""" if isinstance(obj, dict): return {k: _coerce_numbers(v, k) for k, v in obj.items()} if isinstance(obj, list): return [_coerce_numbers(v, key) for v in obj] if key in _NUMERIC_KEYS and isinstance(obj, str): cleaned = obj.replace(",", "").replace("₹", "").replace("Rs.", "").strip() try: n = float(cleaned) return int(n) if n.is_integer() else n except ValueError: return obj return obj def parse_json(raw: str) -> dict: raw = raw.strip() raw = re.sub(r"^```(?:json)?\s*|\s*```$", "", raw) start, end = raw.find("{"), raw.rfind("}") if start == -1 or end == -1: raise ValueError(f"model returned no JSON object: {raw[:200]!r}") raw = raw[start : end + 1] try: data = json.loads(raw) except json.JSONDecodeError: import json_repair data = json_repair.loads(raw) if not isinstance(data, dict): raise ValueError(f"model returned unrepairable JSON: {raw[:200]!r}") return _coerce_numbers(data) def pdf_to_text(pdf_bytes: bytes) -> str: with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: text = "\n\n".join(page.get_text("text", sort=True) for page in doc).strip() if len(text) < 50: raise NoTextLayerError( "PDF contains no usable text layer (scanned image?). OCR is not supported." ) return text # ---------- poller ---------- def _base() -> str: return os.environ["WORKER_URL"].rstrip("/") def _headers(): return {"x-internal-key": os.environ["INTERNAL_KEY"]} def process_one(client: httpx.Client) -> bool: """Claim and process one job. Returns False when the queue is empty.""" resp = client.post(f"{_base()}/internal/claim", headers=_headers(), timeout=30) if resp.status_code == 204: return False resp.raise_for_status() job = resp.json() job_id = job["id"] started = time.time() state["current_job"] = job_id log.info("Processing job %s (%s)", job_id, job.get("filename")) try: pdf_resp = client.get(f"{_base()}/internal/pdf/{job_id}", headers=_headers(), timeout=120) if pdf_resp.status_code == 404: raise FileNotFoundError("PDF not found in storage (expired after 24h?)") pdf_resp.raise_for_status() text = pdf_to_text(pdf_resp.content) result = parse_json(llm_generate(text, job.get("company_hint"))) result = merge_result(result, deterministic_fields(text)) payload = {"result": result, "latency_s": round(time.time() - started, 1)} state["processed"] += 1 log.info("Job %s done in %.1fs", job_id, time.time() - started) except (NoTextLayerError, FileNotFoundError, ValueError, json.JSONDecodeError) as e: payload = {"error": str(e), "latency_s": round(time.time() - started, 1)} state["failed"] += 1 log.warning("Job %s failed: %s", job_id, e) except Exception as e: # Transient (GPU quota, network): don't complete the job; the Worker # re-queues it automatically after the stale timeout. state["last_error"] = f"{type(e).__name__}: {e}" state["current_job"] = None log.exception("Transient error on job %s; leaving it for stale re-queue", job_id) time.sleep(60) return True client.post( f"{_base()}/internal/jobs/{job_id}/complete", headers=_headers(), json=payload, timeout=60, ).raise_for_status() state["current_job"] = None return True def poll_loop(): log.info("Poller started against %s", _base()) with httpx.Client() as client: while True: try: state["last_poll"] = int(time.time()) if process_one(client): continue # drain queue without sleeping except Exception as e: state["last_error"] = f"{type(e).__name__}: {e}" log.exception("Poll loop error; backing off") time.sleep(15) continue time.sleep(POLL_INTERVAL) if os.environ.get("WORKER_URL") and os.environ.get("INTERNAL_KEY"): threading.Thread(target=poll_loop, daemon=True, name="poller").start() else: log.error("WORKER_URL / INTERNAL_KEY not set; poller not started") # ---------- minimal public UI (status only, no data) ---------- def status(): s = dict(state) s["uptime_s"] = int(time.time()) - s.pop("started_at") if s["last_poll"]: s["seconds_since_poll"] = int(time.time()) - s["last_poll"] return s with gr.Blocks(title="invoice-processor worker") as demo: gr.Markdown( "# 🧾 Invoice processor — model worker\n" "Private job worker for its owner's invoice pipeline. No data is accessible here.\n" "Powered by [LiquidAI/LFM2-1.2B-Extract](https://huggingface.co/LiquidAI/LFM2-1.2B-Extract)." ) out = gr.JSON(label="worker status") demo.load(status, outputs=out) gr.Timer(30).tick(status, outputs=out) demo.launch()