Spaces:
Running
on
Zero
Running
on
Zero
File size: 15,100 Bytes
98482ce 55627c9 98482ce |
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 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 |
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
import os
import base64
import tempfile
import json
from PIL import Image
import io
from src.parsers.parser_interface import DocumentParser
from src.parsers.parser_registry import ParserRegistry
# Import the Mistral AI client
try:
from mistralai import Mistral
MISTRAL_AVAILABLE = True
except ImportError:
MISTRAL_AVAILABLE = False
# Load API key from environment variable
api_key = os.getenv("MISTRAL_API_KEY")
# Check if API key is available and print a message if not
if not api_key:
print("Warning: MISTRAL_API_KEY environment variable not found. Mistral OCR parser may not work.")
class MistralOcrParser(DocumentParser):
"""Parser that uses Mistral OCR to convert documents to markdown."""
@classmethod
def get_name(cls) -> str:
return "Mistral OCR (pdf, jpg, png)"
@classmethod
def get_supported_ocr_methods(cls) -> List[Dict[str, Any]]:
return [
{
"id": "ocr",
"name": "OCR Only",
"default_params": {}
},
{
"id": "understand",
"name": "Document Understanding",
"default_params": {}
}
]
@classmethod
def get_description(cls) -> str:
return "Mistral OCR parser for extracting text from documents and images with optional document understanding"
def encode_image(self, image_path):
"""Encode the image to base64."""
try:
with open(image_path, "rb") as image_file:
return base64.b64encode(image_file.read()).decode('utf-8')
except FileNotFoundError:
print(f"Error: The file {image_path} was not found.")
return None
except Exception as e:
print(f"Error: {e}")
return None
def parse(self, file_path: Union[str, Path], ocr_method: Optional[str] = None, **kwargs) -> str:
"""Parse a document using Mistral OCR."""
if not MISTRAL_AVAILABLE:
raise ImportError(
"The Mistral AI client is not installed. "
"Please install it with 'pip install mistralai'."
)
# Use the globally loaded API key
if not api_key:
raise ValueError(
"MISTRAL_API_KEY environment variable is not set. "
"Please set it to your Mistral API key."
)
# Check the OCR method
use_document_understanding = ocr_method == "understand"
try:
# Initialize the Mistral client
client = Mistral(api_key=api_key)
# Determine file type based on extension
file_path = Path(file_path)
file_extension = file_path.suffix.lower()
# Process the document with OCR
if use_document_understanding:
# Use document understanding via chat API for enhanced extraction
return self._extract_with_document_understanding(client, file_path, file_extension)
else:
# Use regular OCR for basic text extraction
return self._extract_with_ocr(client, file_path, file_extension)
except Exception as e:
error_message = f"Error parsing document with Mistral OCR: {str(e)}"
print(error_message)
return f"# Error\n\n{error_message}\n\nPlease check your API key and try again."
def _extract_with_ocr(self, client, file_path, file_extension):
"""Extract document content using basic OCR."""
try:
# Process according to file type
if file_extension in ['.pdf']:
# For PDFs, we need to upload the file to the Mistral API first
try:
# Upload the file to Mistral API
uploaded_pdf = client.files.upload(
file={
"file_name": file_path.name,
"content": open(file_path, "rb"),
},
purpose="ocr"
)
# Get signed URL for the file
signed_url = client.files.get_signed_url(file_id=uploaded_pdf.id)
# Use the signed URL for OCR processing
ocr_response = client.ocr.process(
model="mistral-ocr-latest",
document={
"type": "document_url",
"document_url": signed_url.url
}
)
except Exception as e:
# If file upload fails, try to use a direct URL method with base64
print(f"Failed to upload PDF, trying alternate method: {str(e)}")
base64_pdf = self.encode_image(file_path)
if base64_pdf:
ocr_response = client.ocr.process(
model="mistral-ocr-latest",
document={
"type": "image_url",
"image_url": f"data:application/pdf;base64,{base64_pdf}"
}
)
else:
return "# Error\n\nFailed to process PDF document."
else:
# For images (jpg, png, etc.), use image_url with base64
base64_image = self.encode_image(file_path)
if not base64_image:
return "# Error\n\nFailed to encode the image."
mime_type = self._get_mime_type(file_extension)
ocr_response = client.ocr.process(
model="mistral-ocr-latest",
document={
"type": "image_url",
"image_url": f"data:{mime_type};base64,{base64_image}"
}
)
# Process the OCR response
# The Mistral OCR response is structured with pages that contain text content
markdown_text = ""
# Check if the response contains pages
if hasattr(ocr_response, 'pages') and ocr_response.pages:
for page in ocr_response.pages:
# Add page number as heading
page_num = page.index if hasattr(page, 'index') else "Unknown"
markdown_text += f"## Page {page_num}\n\n"
# Add text content if available
if hasattr(page, 'text'):
markdown_text += page.text + "\n\n"
# Or markdown content if that's how it's structured
elif hasattr(page, 'markdown'):
markdown_text += page.markdown + "\n\n"
# Add any extracted tables with markdown formatting
if hasattr(page, 'tables') and page.tables:
for i, table in enumerate(page.tables):
markdown_text += f"### Table {i+1}\n\n"
if hasattr(table, 'markdown'):
markdown_text += table.markdown + "\n\n"
elif hasattr(table, 'data'):
# Convert table data to markdown format
markdown_text += self._convert_table_to_markdown(table.data) + "\n\n"
# If no markdown was generated, check for raw content
if not markdown_text and hasattr(ocr_response, 'content'):
markdown_text = ocr_response.content
# If still no content, try to access any available data
if not markdown_text:
# Try to get a JSON representation to extract data
try:
response_dict = ocr_response.to_dict() if hasattr(ocr_response, 'to_dict') else ocr_response.__dict__
markdown_text = "# Extracted Content\n\n"
# Look for content or text in the response dictionary
if 'content' in response_dict:
markdown_text += response_dict['content']
elif 'text' in response_dict:
markdown_text += response_dict['text']
elif 'pages' in response_dict:
for page in response_dict['pages']:
if 'text' in page:
markdown_text += page['text'] + "\n\n"
else:
# Just dump what we got as JSON
markdown_text += f"```json\n{json.dumps(response_dict, indent=2)}\n```"
except Exception as e:
markdown_text = f"# Error Processing Response\n\nCould not process the OCR response: {str(e)}"
# If we still have no content, return an error
if not markdown_text:
return "# Error\n\nNo text was extracted from the document."
return f"# Document Content\n\n{markdown_text}"
except Exception as e:
return f"# OCR Extraction Error\n\n{str(e)}"
def _extract_with_document_understanding(self, client, file_path, file_extension):
"""Extract and understand document content using chat completion."""
try:
# For PDFs and images, we'll use Mistral's document understanding capability
if file_extension in ['.pdf']:
# Upload PDF first
try:
# Upload the file
uploaded_pdf = client.files.upload(
file={
"file_name": file_path.name,
"content": open(file_path, "rb"),
},
purpose="ocr"
)
# Get the signed URL
signed_url = client.files.get_signed_url(file_id=uploaded_pdf.id)
# Send to chat completion API with document understanding prompt
chat_response = client.chat.complete(
model="mistral-large-latest",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": "Convert this document to well-formatted markdown. Preserve all important content, structure, headings, lists, and tables. Include brief descriptions of any images."
},
{
"type": "document_url",
"document_url": signed_url.url
}
]
}
]
)
# Get the markdown result
return chat_response.choices[0].message.content
except Exception as e:
# Fall back to OCR if document understanding fails
print(f"Document understanding failed, falling back to OCR: {str(e)}")
return self._extract_with_ocr(client, file_path, file_extension)
else:
# For images, encode to base64
base64_image = self.encode_image(file_path)
if not base64_image:
return "# Error\n\nFailed to encode the image."
mime_type = self._get_mime_type(file_extension)
# Use the chat API with the image for document understanding
chat_response = client.chat.complete(
model="mistral-large-latest",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": "Extract all text from this image and convert it to well-formatted markdown. Preserve the structure and layout as much as possible."
},
{
"type": "image_url",
"image_url": {
"url": f"data:{mime_type};base64,{base64_image}"
}
}
]
}
]
)
# Get the markdown result
return chat_response.choices[0].message.content
except Exception as e:
return f"# Document Understanding Error\n\n{str(e)}\n\nFalling back to OCR method."
def _get_mime_type(self, file_extension: str) -> str:
"""Get the MIME type for a file extension."""
mime_types = {
".pdf": "application/pdf",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".png": "image/png",
".gif": "image/gif",
".bmp": "image/bmp",
".tiff": "image/tiff",
".tif": "image/tiff",
}
return mime_types.get(file_extension, "application/octet-stream")
def _convert_table_to_markdown(self, table_data) -> str:
"""Convert a table data structure to markdown format."""
if not table_data or not isinstance(table_data, list):
return ""
# Create markdown table
markdown = ""
# Add header row
if table_data and isinstance(table_data[0], list):
header = table_data[0]
markdown += "| " + " | ".join(str(cell) for cell in header) + " |\n"
# Add separator row
markdown += "| " + " | ".join(["---"] * len(header)) + " |\n"
# Add data rows
for row in table_data[1:]:
markdown += "| " + " | ".join(str(cell) for cell in row) + " |\n"
return markdown
# Register the parser with the registry
if MISTRAL_AVAILABLE:
ParserRegistry.register(MistralOcrParser)
else:
print("Mistral OCR parser not registered: mistralai package not installed") |