55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
import ollama
|
|
import json
|
|
import base64
|
|
|
|
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 = """
|
|
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)
|
|
|
|
IMPORTANT:
|
|
- Return ONLY the JSON. No markdown formatting, no explanations.
|
|
- If a field is not found, 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 ""}
|
|
|