Add BillStructAI demo

app.py

@@ -0,0 +1,153 @@
+"""
+BillStructAI demo — paste noisy OCR text from an invoice/receipt and get
+back structured JSON, using a Qwen2.5-1.5B-Instruct base model fine-tuned
+with QLoRA for this task.
+
+Deploy: push this file + requirements.txt to a Hugging Face Space
+(Gradio SDK). Set ADAPTER_REPO below to your uploaded LoRA adapter repo.
+"""
+
+import json
+import re
+
+import gradio as gr
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+BASE_MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+ADAPTER_REPO = "trishpurkait/billstructai-qwen-lora"  # <-- update after uploading adapter
+
+SYSTEM_PROMPT = """
+You are a precise invoice information extraction assistant.
+Your task is to extract structured invoice data from noisy OCR text.
+Return only valid JSON.
+Do not explain.
+Do not add markdown.
+Use null for missing fields.
+Do not hallucinate values that are not present in the OCR text.
+""".strip()
+
+EXAMPLE_OCR = """RECEIPT
+ORLANDO INTERNATIONAL AIRPORT
+2001207734
+EXIT: SEQUENCE /LINE/CASHIER/ DATE / TIME / FEE /COST LD.
+3493 11 108 18DE 0008 012.00 FLFRC927 R01 67212 2. 14DE 0817"""
+
+_model = None
+_tokenizer = None
+
+
+def load_model():
+    """Lazily load base model + LoRA adapter (4-bit) on first request."""
+    global _model, _tokenizer
+    if _model is not None:
+        return _model, _tokenizer
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_NAME,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
+    model.eval()
+
+    _model, _tokenizer = model, tokenizer
+    return model, tokenizer
+
+
+def extract_json_from_text(text: str):
+    """Best-effort JSON extraction from model output (handles markdown fences)."""
+    if not text:
+        return None
+    text = text.strip()
+    try:
+        return json.loads(text)
+    except Exception:
+        pass
+    text = text.replace("```json", "").replace("```", "").strip()
+    try:
+        return json.loads(text)
+    except Exception:
+        pass
+    match = re.search(r"\{.*\}", text, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(0))
+        except Exception:
+            return None
+    return None
+
+
+def run_extraction(ocr_text: str) -> str:
+    if not ocr_text or not ocr_text.strip():
+        return "Paste some OCR text first."
+
+    model, tokenizer = load_model()
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": f"Extract invoice information from the OCR text below and "
+            f"return only valid JSON.\n\nOCR_TEXT:\n{ocr_text.strip()}",
+        },
+    ]
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=768,
+            do_sample=False,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+
+    generated = tokenizer.decode(
+        output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
+    )
+
+    parsed = extract_json_from_text(generated)
+    if parsed is not None:
+        return json.dumps(parsed, indent=2, ensure_ascii=False)
+    return generated  # fall back to raw output if parsing fails
+
+
+with gr.Blocks(title="BillStructAI — Invoice OCR to JSON") as demo:
+    gr.Markdown(
+        "# BillStructAI\n"
+        "Paste noisy OCR text from an invoice or receipt below. "
+        "A Qwen2.5-1.5B-Instruct model fine-tuned with QLoRA extracts it into "
+        "structured JSON. [See the full writeup and evaluation results](https://github.com/your-username/billstructai)."
+    )
+    with gr.Row():
+        with gr.Column():
+            ocr_input = gr.Textbox(
+                label="OCR text",
+                lines=12,
+                value=EXAMPLE_OCR,
+                placeholder="Paste raw OCR text here...",
+            )
+            submit_btn = gr.Button("Extract", variant="primary")
+        with gr.Column():
+            json_output = gr.Code(label="Extracted JSON", language="json", lines=20)
+
+    submit_btn.click(fn=run_extraction, inputs=ocr_input, outputs=json_output)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio>=4.0.0
+torch
+transformers
+accelerate
+bitsandbytes
+peft