Spaces:
Running on Zero
Running on Zero
File size: 6,012 Bytes
8e798a2 0de8787 8e798a2 0de8787 8e798a2 7b48434 8e798a2 | 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 | import time
import logging
from typing import Any, Dict
from .config import settings
from .ocr_service import ocr_service, OCRService
from .nvidia_client import nvidia_client, NVIDIAClient
from .prompts import DOCUMENT_INTELLIGENCE_SYSTEM_PROMPT, build_analysis_prompt
from .schemas import DocumentIntelligenceResult
logger = logging.getLogger("document_intelligence.analyzer")
class DocumentAnalyzer:
"""Enterprise document intelligence orchestrator combining OCR and NVIDIA LLM."""
def __init__(
self,
ocr: OCRService = ocr_service,
llm_client: NVIDIAClient = nvidia_client
):
self.ocr = ocr
self.llm_client = llm_client
def analyze(self, image_input: Any) -> DocumentIntelligenceResult:
"""
End-to-end pipeline:
Image -> Preprocessing -> OCR Extraction -> Prompt Formatting -> NVIDIA LLM -> Pydantic Schema Validation.
"""
overall_start = time.time()
logger.info("Starting document analysis pipeline.")
# 1. OCR Text Extraction
try:
ocr_result = self.ocr.extract_text(image_input)
ocr_text = ocr_result.get("full_text", "").strip()
ocr_hints = ocr_result.get("spatial_hints", "")
ocr_latency = ocr_result.get("latency_seconds", 0.0)
except Exception as ocr_err:
logger.error("OCR extract error: %s", ocr_err, exc_info=True)
return DocumentIntelligenceResult(
document_type="unknown",
extracted_text=f"[OCR Processing Error: {str(ocr_err)}]",
entities=[],
tables=[],
components=[],
relationships=[],
summary=f"Failed to process image through OCR: {str(ocr_err)}",
confidence_notes=[f"OCR error details: {str(ocr_err)}", f"Input received type: {type(image_input)}"],
metadata={"error": str(ocr_err), "input_type": str(type(image_input))}
)
if not ocr_text:
logger.warning("No text detected by OCR engine.")
return DocumentIntelligenceResult(
document_type="unknown",
extracted_text="[No readable text detected in the provided image]",
entities=[],
tables=[],
components=[],
relationships=[],
summary="The OCR engine could not identify any readable text or characters in the uploaded image.",
confidence_notes=[
"Low quality image or no text present.",
f"Engine attempted: {ocr_result.get('engine', 'unknown')}",
f"Image dimensions: {ocr_result.get('image_size', 'unknown')}"
],
metadata={
"ocr_engine": ocr_result.get("engine", "unknown"),
"ocr_latency_seconds": ocr_latency,
"total_latency_seconds": round(time.time() - overall_start, 4)
}
)
# 2. Build Prompt
user_prompt = build_analysis_prompt(ocr_text, ocr_hints)
# 3. LLM Structured Extraction
try:
llm_response = self.llm_client.generate_document_intelligence(
system_prompt=DOCUMENT_INTELLIGENCE_SYSTEM_PROMPT,
user_prompt=user_prompt
)
parsed_data = llm_response.get("parsed_data", {})
llm_latency = llm_response.get("latency_seconds", 0.0)
model_name = llm_response.get("model", settings.NVIDIA_MODEL)
except Exception as llm_err:
logger.warning("LLM call failed: %s. Returning OCR-only response.", llm_err)
total_latency = time.time() - overall_start
return DocumentIntelligenceResult(
document_type="raw_document",
extracted_text=ocr_text,
entities=[],
tables=[],
components=[],
relationships=[],
summary=f"OCR text extraction succeeded ({len(ocr_result.get('lines', []))} lines detected). LLM processing skipped due to: {str(llm_err)}",
confidence_notes=[
f"OCR Engine [{ocr_result.get('engine', 'unknown')}] extracted {len(ocr_result.get('lines', []))} lines in {ocr_latency:.3f}s.",
f"NVIDIA LLM Error: {str(llm_err)}. Please set a valid NVIDIA_API_KEY in Space Secrets."
],
metadata={
"ocr_engine": ocr_result.get("engine", "unknown"),
"ocr_lines_detected": len(ocr_result.get("lines", [])),
"ocr_latency_seconds": ocr_latency,
"llm_error": str(llm_err),
"total_latency_seconds": round(total_latency, 4)
}
)
# 4. Fallback values if LLM skipped some keys
if "extracted_text" not in parsed_data or not parsed_data["extracted_text"]:
parsed_data["extracted_text"] = ocr_text
# 5. Attach execution metadata
total_latency = time.time() - overall_start
parsed_data["metadata"] = {
"ocr_engine": ocr_result.get("engine", "unknown"),
"ocr_lines_detected": len(ocr_result.get("lines", [])),
"ocr_latency_seconds": ocr_latency,
"llm_model": model_name,
"llm_latency_seconds": llm_latency,
"total_latency_seconds": round(total_latency, 4),
}
# 6. Pydantic validation
result = DocumentIntelligenceResult(**parsed_data)
logger.info(
"Document analysis complete. Type: %s, Entities: %d, Tables: %d, Components: %d in %.3f s",
result.document_type,
len(result.entities),
len(result.tables),
len(result.components),
total_latency
)
return result
# Global analyzer instance
document_analyzer = DocumentAnalyzer()
|