Spaces:
Sleeping
Sleeping
File size: 6,908 Bytes
c1b1f3a | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | import os
import re
import json
import requests
from typing import Dict, Any, List, Tuple, Optional
from ..config import get_secret
from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS
from .hybrid_extractor import HybridFieldExtractor
class CohereFieldExtractor:
"""
AI-Powered Document Field Extractor using Cohere API (command-r-08-2024).
Features:
- Deep semantic extraction across Invoices, Receipts, and Application Forms
- Structured JSON output conforming to CLASS_FIELD_SCHEMAS
- Dynamic alignment of extracted entities back to OCRWord bounding boxes
- Automatic graceful fallback to HybridFieldExtractor if API key is missing or offline
"""
API_URL = "https://api.cohere.com/v2/chat"
MODEL = "command-r-08-2024"
def __init__(self, doc_class: str, api_key: Optional[str] = None):
self.doc_class = doc_class
self.schema_fields = CLASS_FIELD_SCHEMAS.get(doc_class, [])
self.heuristic_fallback = HybridFieldExtractor(doc_class)
# Load API key
if not api_key:
api_key = get_secret("COHERE_API", "COHERE_API_KEY", default="")
self.api_key = api_key.strip().strip("'\"")
def _call_cohere_extraction(self, full_text: str) -> Dict[str, Any]:
"""Calls Cohere v2 Chat API to extract structured fields into JSON."""
if not self.api_key:
raise RuntimeError("COHERE_API key not found in environment or .env file.")
system_prompt = (
f"You are an expert enterprise document information extraction AI. "
f"Your task is to accurately extract key entity fields from OCR text of a '{self.doc_class}' document. "
f"Required fields to extract: {self.schema_fields}.\n"
f"Rules:\n"
f"1. Extract exact, truthful values as written in the document.\n"
f"2. For amounts, extract the total numeric amount due/paid (e.g. '15,726.56').\n"
f"3. For dates, extract the document date.\n"
f"4. If a field is not present in the document, assign its value to null.\n"
f"5. Output must be a strictly valid JSON object matching the requested field keys."
)
user_content = f"Document Type: {self.doc_class.replace('_', ' ').title()}\nDocument OCR Text:\n{full_text}"
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
payload = {
"model": self.MODEL,
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_content}
],
"response_format": {"type": "json_object"},
"temperature": 0.0
}
response = requests.post(self.API_URL, headers=headers, json=payload, timeout=20)
response.raise_for_status()
data = response.json()
raw_content = data.get("message", {}).get("content", [])
if raw_content and isinstance(raw_content, list):
text_block = raw_content[0].get("text", "{}")
extracted_json = json.loads(text_block)
return extracted_json
return {}
def _find_bbox_for_value(self, value_str: str, all_words: List[OCRWord]) -> Optional[List[int]]:
"""
Maps an extracted text string back to spatial OCRWord bounding boxes.
Handles both multi-word phrases (e.g. company names) and single tokens (e.g. IDs, numbers).
"""
if not value_str or not all_words:
return None
val_clean = str(value_str).strip()
tokens = [t for t in re.split(r'[\s,.:;/\-]+', val_clean) if len(t) >= 2]
if not tokens:
tokens = [val_clean]
matching_words: List[OCRWord] = []
# 1. Exact or partial match for tokens
for token in tokens:
t_lower = token.lower()
best_w = None
for w in all_words:
w_lower = w.text.lower().strip()
if t_lower == w_lower or (len(t_lower) > 3 and t_lower in w_lower):
best_w = w
break
if best_w and best_w not in matching_words:
matching_words.append(best_w)
# 2. If matched words found, compute union bounding box
if matching_words:
x0 = min(w.bbox[0] for w in matching_words)
y0 = min(w.bbox[1] for w in matching_words)
x1 = max(w.bbox[2] for w in matching_words)
y1 = max(w.bbox[3] for w in matching_words)
return [int(x0), int(y0), int(x1), int(y1)]
return None
def extract_with_details(self, record: DocumentRecord) -> Tuple[Dict[str, Any], Dict[str, List[int]], Dict[str, float]]:
"""
Extracts structured fields using Cohere AI, maps spatial bounding boxes,
and uses HybridFieldExtractor for any missing fields or fallback.
"""
words = record.ocr_words
full_text = " ".join([w.text for w in words]).strip()
extracted_fields = {f: None for f in self.schema_fields}
bboxes = {f: None for f in self.schema_fields}
confidences = {f: 0.0 for f in self.schema_fields}
ai_succeeded = False
if self.api_key and full_text:
try:
ai_results = self._call_cohere_extraction(full_text)
for f in self.schema_fields:
val = ai_results.get(f)
if val is not None and str(val).strip() and str(val).lower() not in ["null", "none", "not detected"]:
extracted_fields[f] = str(val).strip()
confidences[f] = 0.95
# Map spatial bounding box
bbox = self._find_bbox_for_value(str(val), words)
if bbox:
bboxes[f] = bbox
ai_succeeded = True
except Exception as e:
print(f"[WARN] Cohere AI extraction error: {e}. Falling back to heuristic extractor.")
# If AI failed completely or certain fields were not found, use heuristic extractor to fill gaps
heuristic_fields, heuristic_bboxes, heuristic_confs = self.heuristic_fallback.extract_with_details(record)
for f in self.schema_fields:
if extracted_fields.get(f) is None and heuristic_fields.get(f) is not None:
extracted_fields[f] = heuristic_fields[f]
bboxes[f] = heuristic_bboxes.get(f)
confidences[f] = heuristic_confs.get(f, 0.80)
elif bboxes.get(f) is None and heuristic_bboxes.get(f) is not None:
# Borrow heuristic bbox if AI found the value but token matching was fuzzy
bboxes[f] = heuristic_bboxes.get(f)
return extracted_fields, bboxes, confidences
|