added azure document intelligence ocr

app.py

@@ -12,8 +12,11 @@ from policy_analyser.analyse import Health
 if 'GPT_KEY' not in os.environ or os.environ.get('GPT_KEY') in [None, '']:
     os.environ['GPT_KEY'] = st.secrets['GPT_KEY']
 
+if 'AZURE_LAYOUT_KEY' not in os.environ.get('AZURE_LAYOUT_KEY') or os.environ.get('AZURE_LAYOUT_KEY') in [None, '']:
+    os.environ['AZURE_LAYOUT_KEY'] = st.secrets['AZURE_LAYOUT_KEY']
+
 if 'health_analyser' not in st.session_state:
-    st.session_state.health_analyser = Health()
+    st.session_state.health_analyser = Health(ocr_engine = 'azure/layout')
 
 def markdown_table_to_json(markdown):
     lines = markdown.strip().split("\n")

policy_analyser/__init__.py

@@ -31,6 +31,9 @@ GPT_ENGINE = 'o3-mini'
 GPT_KEY = os.environ.get('GPT_KEY', '')
 GPT_VERSION = '2024-12-01-preview'
 GPT_API_BASE = 'https://ai-ackods910341544474.openai.azure.com/'
+AZURE_LAYOUT_ENDPOINT = 'https://acko-document-intelligence.cognitiveservices.azure.com/'
+AZURE_LAYOUT_KEY = os.environ.get('AZURE_LAYOUT_KEY', '')
+AZURE_LAYOUT_MODEL = 'prebuilt-layout'
 
 # EXTRACTION_PROMPT = open(os.path.join(PROMPTS_DIR, 'extraction.txt')).read()
 # entities = json.load(open(os.path.join(DATA_DIR, 'policy_analyser_entities.json')))

policy_analyser/analyse.py

@@ -9,7 +9,7 @@ from time import time
 from datetime import datetime
 
 from policy_analyser import PROMPTS_DIR, DATA_DIR
-from policy_analyser.ocr import PyMuPDF4LLMOCR
+from policy_analyser.ocr import PyMuPDF4LLMOCR, AzureDocumentIntelligenceOCR
 from policy_analyser.llm import call_openai
 from policy_analyser.utils import markdown_table_to_json
 
@@ -17,6 +17,8 @@ class LOB:
     def __init__(self, ocr_engine = 'open-source/pymupdf4llm'):
         if ocr_engine == 'open-source/pymupdf4llm':
             self.engine = PyMuPDF4LLMOCR()
+        elif ocr_engine == 'azure/layout':
+            self.engine = AzureDocumentIntelligenceOCR()
         self.file_type = 'pdf'
         with open(os.path.join(PROMPTS_DIR, 'analysis.txt'), 'r') as f:
             self.analysis_prompt = f.read()

policy_analyser/ocr.py

@@ -5,6 +5,11 @@
 
 # Imports
 import pymupdf4llm, pymupdf
+from azure.core.credentials import AzureKeyCredential
+from azure.ai.documentintelligence import DocumentIntelligenceClient
+from azure.ai.documentintelligence.models import AnalyzeDocumentRequest, ContentFormat, AnalyzeResult
+
+from policy_analyser import AZURE_LAYOUT_ENDPOINT, AZURE_LAYOUT_KEY, AZURE_LAYOUT_MODEL
 
 class PyMuPDF4LLMOCR:
     def __init__(self):
@@ -19,6 +24,33 @@ class PyMuPDF4LLMOCR:
         response = pymupdf4llm.to_markdown(document)
         return response, None
 
+class AzureDocumentIntelligenceOCR:
+    def __init__(self):
+        self.engine = 'azure/layout'
+        self.client = None
+
+    def _authenticate(self):
+        if self.client is None:
+            try:
+                self.client = DocumentIntelligenceClient(
+                    endpoint = AZURE_LAYOUT_ENDPOINT,
+                    credential = AzureKeyCredential(AZURE_LAYOUT_KEY)
+                )
+            except Exception as e:
+                self.client = None
+
+    def __call__(self, file_bytes):
+        if self.client is not None:
+            poller = self.client.begin_analyze_document(
+                AZURE_LAYOUT_MODEL,
+                AnalyzeDocumentRequest(bytes_source = file_bytes),
+                output_content_format = ContentFormat.MARKDOWN
+            )
+            result = poller.result()
+            return result.content, None
+        else:
+            print('Client is not authenticated or reachable')
+
 if __name__ == '__main__':
     import sys
     filepath = sys.argv[1]