Spaces:
Sleeping
Sleeping
| 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 | |