117 lines
3.1 KiB
Python
117 lines
3.1 KiB
Python
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 = f"""
|
|
You are an expert data extraction assistant.
|
|
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, no prefix.
|
|
- If a field is not found or data is not available, use null.
|
|
"""
|
|
|
|
messages = [{'role': 'user', 'content': prompt}]
|
|
model = 'gemma:2b'
|
|
|
|
if model_type == 'vision':
|
|
if not image_path:
|
|
return {"error": "Image path required for vision mode"}
|
|
|
|
# Qwen-VL handles images passed in the message
|
|
model = 'qwen2.5vl:7b' # Using the installed model ID
|
|
messages[0]['images'] = [image_path]
|
|
messages[0]['content'] = "Analyze this image. " + prompt
|
|
else:
|
|
# Text Mode
|
|
if not text:
|
|
return {"error": "Text required for text mode"}
|
|
messages[0]['content'] += f"\n\n---\n{text}\n---"
|
|
|
|
try:
|
|
response = ollama.chat(model=model, messages=messages)
|
|
content = response['message']['content']
|
|
|
|
# Clean up markdown
|
|
content = content.replace("```json", "").replace("```", "").strip()
|
|
|
|
return json.loads(content)
|
|
|
|
except Exception as e:
|
|
print(f"LLM Extraction Error ({model_type}): {e}")
|
|
return {"error": str(e), "raw_output": content if 'content' in locals() else ""}
|
|
|