Spaces:
Sleeping
Sleeping
File size: 3,693 Bytes
d24a0cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import os
import base64
import json
from typing import Any, Dict
import httpx
# Get your OpenRouter API key from env (you'll set this in Hugging Face later)
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL_NAME = "qwen/qwen3-vl-235b-a22b-instruct"
def _file_to_image_block(file_bytes: bytes, content_type: str) -> Dict[str, Any]:
"""
Encode the file as a data URL to feed into the multimodal model.
For demo purposes we treat PDFs and images the same way here.
"""
b64 = base64.b64encode(file_bytes).decode("utf-8")
return {
"type": "input_image",
"image_url": f"data:{content_type};base64,{b64}",
}
async def extract_fields_from_document(
file_bytes: bytes,
content_type: str,
filename: str,
) -> Dict[str, Any]:
"""
Call OpenRouter with Qwen3-VL and return parsed JSON with fields.
We instruct the model to return JSON only.
"""
if not OPENROUTER_API_KEY:
raise RuntimeError("OPENROUTER_API_KEY environment variable is not set")
image_block = _file_to_image_block(file_bytes, content_type)
system_prompt = (
"You are a document extraction engine. "
"You analyze invoices, receipts, contracts, reports and similar documents, "
"and output structured JSON only (no explanations or comments)."
)
user_prompt = (
"Extract important key-value pairs from the document and respond with JSON only.\n"
"Use this shape:\n"
"{\n"
' \"doc_type\": \"invoice | receipt | contract | report | other\",\n'
' \"confidence\": number between 0 and 100,\n'
' \"fields\": {\n'
' \"invoice_number\": \"...\",\n'
' \"date\": \"...\",\n'
' \"due_date\": \"...\",\n'
' \"total_amount\": \"...\",\n'
' \"currency\": \"...\",\n'
' \"vendor_name\": \"...\",\n'
' \"line_items\": [\n'
' {\"description\": \"...\", \"quantity\": \"...\", \"unit_price\": \"...\", \"line_total\": \"...\"}\n'
' ],\n'
' \"other_field\": \"...\"\n'
" }\n"
"}\n"
"If fields are missing or not applicable, simply omit them."
)
payload: Dict[str, Any] = {
"model": MODEL_NAME,
"messages": [
{
"role": "system",
"content": [{"type": "text", "text": system_prompt}],
},
{
"role": "user",
"content": [
{"type": "text", "text": user_prompt},
image_block,
],
},
],
"max_tokens": 2048,
}
headers = {
"Authorization": f"Bearer {OPENROUTER_API_KEY}",
"Content-Type": "application/json",
# Optional attribution headers
"HTTP-Referer": os.environ.get(
"APP_URL",
"https://huggingface.co/spaces/your-space",
),
"X-Title": "Document Capture Demo",
}
async with httpx.AsyncClient(timeout=120) as client:
resp = await client.post(OPENROUTER_BASE_URL, headers=headers, json=payload)
resp.raise_for_status()
data = resp.json()
# OpenRouter returns choices[0].message.content
content = data["choices"][0]["message"]["content"]
# content may be a string or a list of content blocks
if isinstance(content, list):
text = "".join(part.get("text", "") for part in content if part.get("type") == "text")
else:
text = content
# Try to parse JSON from the model output
return json.loads(text)
|