""" Author : Janarddan Sarkar file_name : mistral_ocr_st.py date : 10-03-2025 description : Fixed JSON parsing by using manual JSON extraction with response_format=json_object. """ import os import json import base64 import time import streamlit as st from mistralai import Mistral from dotenv import find_dotenv, load_dotenv from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk from mistralai.models import OCRResponse from mistralai.models.sdkerror import SDKError from enum import Enum from pydantic import BaseModel import pycountry # Load environment variables load_dotenv(find_dotenv()) api_key = os.environ.get("MISTRAL_API_KEY") client = Mistral(api_key=api_key) # Define Language Enum (保留仅供其他用途，但不在模型中使用) languages = {lang.alpha_2: lang.name for lang in pycountry.languages if hasattr(lang, 'alpha_2')} class LanguageMeta(Enum.__class__): def __new__(metacls, cls, bases, classdict): for code, name in languages.items(): classdict[name.upper().replace(' ', '_')] = name return super().__new__(metacls, cls, bases, classdict) class Language(Enum, metaclass=LanguageMeta): pass # 不再使用 Pydantic 模型作为 response_format，仅定义结构供参考 class StructuredOCR(BaseModel): file_name: str topics: list[str] languages: list[str] ocr_contents: dict def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str: for img_name, base64_str in images_dict.items(): markdown_str = markdown_str.replace(f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})") return markdown_str def get_combined_markdown(ocr_response: OCRResponse) -> str: markdowns: list[str] = [] for page in ocr_response.pages: image_data = {img.id: img.image_base64 for img in page.images} markdowns.append(replace_images_in_markdown(page.markdown, image_data)) return "\n\n".join(markdowns) # ---------- 重试装饰器 ---------- def retry_on_rate_limit(max_retries=5, initial_delay=1): def decorator(func): def wrapper(*args, **kwargs): for attempt in range(max_retries): try: return func(*args, **kwargs) except SDKError as e: if "429" in str(e) and attempt < max_retries - 1: delay = initial_delay * (2 ** attempt) st.warning(f"Rate limit hit. Retrying in {delay}s... (attempt {attempt+1}/{max_retries})") time.sleep(delay) else: raise return None return wrapper return decorator @retry_on_rate_limit(max_retries=5, initial_delay=1) def ocr_process_with_retry(**kwargs): return client.ocr.process(**kwargs) @retry_on_rate_limit(max_retries=3, initial_delay=1) def chat_complete_with_retry(**kwargs): return client.chat.complete(**kwargs) # ---------- 处理函数 ---------- def process_pdf(pdf_bytes, file_name): uploaded_file = client.files.upload( file={"file_name": file_name, "content": pdf_bytes}, purpose="ocr", ) signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1) pdf_response = ocr_process_with_retry( document=DocumentURLChunk(document_url=signed_url.url), model="mistral-ocr-latest", include_image_base64=True, ) if isinstance(pdf_response, dict): pdf_response = OCRResponse(**pdf_response) return pdf_response def process_image(image_bytes, file_name): """Process an image: returns (markdown_with_images, structured_json)""" encoded_image = base64.b64encode(image_bytes).decode() base64_data_url = f"data:image/jpeg;base64,{encoded_image}" # OCR 调用 image_response = ocr_process_with_retry( document=ImageURLChunk(image_url=base64_data_url), model="mistral-ocr-latest", include_image_base64=True ) page = image_response.pages[0] images_dict = {img.id: img.image_base64 for img in page.images if img.image_base64} markdown_with_images = replace_images_in_markdown(page.markdown, images_dict) # 默认 fallback 结构（包含完整的 OCR 文本） default_structured = { "file_name": file_name, "topics": [], "languages": [], "ocr_contents": {"full_text": markdown_with_images} } try: response = chat_complete_with_retry( model="pixtral-12b-latest", messages=[ { "role": "user", "content": [ ImageURLChunk(image_url=base64_data_url), TextChunk( text=( "Based on the image and its OCR markdown below, output a JSON object with the following fields:\n" "- file_name: string, the original file name.\n" "- topics: list of strings, main subjects.\n" "- languages: list of strings, languages detected (e.g., 'English', 'Chinese').\n" "- ocr_contents: dictionary, where keys are section titles or 'full_text', and values are the corresponding text. Include the full OCR text under key 'full_text'.\n\n" f"OCR Markdown:\n{markdown_with_images}\n\n" "Output ONLY valid JSON, no other text." ) ), ], }, ], response_format={"type": "json_object"}, temperature=0, ) content = response.choices[0].message.content parsed_json = json.loads(content) structured_json = { "file_name": parsed_json.get("file_name", file_name), "topics": parsed_json.get("topics", []), "languages": parsed_json.get("languages", []), "ocr_contents": parsed_json.get("ocr_contents", {"full_text": markdown_with_images}) } # 确保类型正确 if not isinstance(structured_json["topics"], list): structured_json["topics"] = [] if not isinstance(structured_json["languages"], list): structured_json["languages"] = [] if not isinstance(structured_json["ocr_contents"], dict): structured_json["ocr_contents"] = {"full_text": markdown_with_images} except Exception as e: st.warning(f"Could not parse structured JSON from model: {e}. Using fallback OCR text as ocr_contents.") structured_json = default_structured return markdown_with_images, structured_json # ---------- Streamlit UI ---------- st.title("OLMOCR") uploaded_file = st.file_uploader("Upload a PDF or Image", type=["pdf", "png", "jpg", "jpeg"]) if uploaded_file: file_type = uploaded_file.type file_bytes = uploaded_file.read() file_name = uploaded_file.name if st.button("Submit"): st.write(f"**Processing file:** {file_name}") if "pdf" in file_type: pdf_response = process_pdf(file_bytes, file_name) st.markdown(get_combined_markdown(pdf_response)) else: markdown_content, structured_data = process_image(file_bytes, file_name) st.markdown(markdown_content) with st.expander("📋 Structured Data (JSON)"): st.json(structured_data)