Changes committed
This commit is contained in:
@@ -2,24 +2,86 @@ import ollama
|
||||
import json
|
||||
import base64
|
||||
|
||||
INVOICE_SCHEMA = {
|
||||
"document_type": None,
|
||||
"invoice_number": None,
|
||||
"invoice_date": None,
|
||||
"due_date": None,
|
||||
"purchase_order_number": None,
|
||||
"vendor": {
|
||||
"name": None,
|
||||
"address": None,
|
||||
"email": None,
|
||||
"phone": None,
|
||||
"gstin": None,
|
||||
"tax_id": None,
|
||||
"website": None
|
||||
},
|
||||
"customer": {
|
||||
"name": None,
|
||||
"address": None,
|
||||
"gstin": None
|
||||
},
|
||||
"amounts": {
|
||||
"subtotal": None,
|
||||
"tax": None,
|
||||
"discount": None,
|
||||
"shipping": None,
|
||||
"round_off": None,
|
||||
"total": None,
|
||||
"amount_paid": None,
|
||||
"balance_due": None,
|
||||
"currency": None
|
||||
},
|
||||
"tax_breakdown": [
|
||||
{
|
||||
"type": None,
|
||||
"rate": None,
|
||||
"amount": None
|
||||
}
|
||||
],
|
||||
"line_items": [
|
||||
{
|
||||
"line_no": None,
|
||||
"description": None,
|
||||
"product_code": None,
|
||||
"hsn_sac": None,
|
||||
"quantity": None,
|
||||
"unit": None,
|
||||
"unit_price": None,
|
||||
"discount": None,
|
||||
"tax_rate": None,
|
||||
"tax_amount": None,
|
||||
"total": None
|
||||
}
|
||||
],
|
||||
"payment_information": {
|
||||
"bank_name": None,
|
||||
"account_number": None,
|
||||
"ifsc": None,
|
||||
"upi_id": None
|
||||
},
|
||||
"metadata": {
|
||||
"pages": None,
|
||||
"ocr_confidence": None,
|
||||
"language": None
|
||||
}
|
||||
}
|
||||
|
||||
def extract_data(text: str = None, image_path: str = None, model_type: str = "text") -> dict:
|
||||
"""
|
||||
Extracts structured data using either Text (Gemma) or Vision (Qwen) models.
|
||||
"""
|
||||
|
||||
prompt = """
|
||||
prompt = f"""
|
||||
You are an expert data extraction assistant.
|
||||
Extract the following fields from the provided document and return them as a SINGLE VALID JSON OBJECT:
|
||||
- invoice_number (string)
|
||||
- date (string)
|
||||
- vendor_name (string)
|
||||
- total_amount (string)
|
||||
- currency (string)
|
||||
- line_items (list of objects with: description, quantity, unit_price, total)
|
||||
Extract every possible detail from the provided document and return it strictly as a SINGLE VALID JSON OBJECT matching the following schema structure:
|
||||
|
||||
{json.dumps(INVOICE_SCHEMA, indent=4)}
|
||||
|
||||
IMPORTANT:
|
||||
- Return ONLY the JSON. No markdown formatting, no explanations.
|
||||
- If a field is not found, use null.
|
||||
- Return ONLY the JSON. No markdown formatting, no explanations, no prefix.
|
||||
- If a field is not found or data is not available, use null.
|
||||
"""
|
||||
|
||||
messages = [{'role': 'user', 'content': prompt}]
|
||||
|
||||
Reference in New Issue
Block a user