import gradio as gr from PIL import Image from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer, TextIteratorStreamer import spaces from threading import Thread from pdf2image import convert_from_path import os import tempfile import base64 from io import BytesIO import time # --- Model Loading --- # Load the model, processor, and tokenizer once when the app starts. model_path = "nanonets/Nanonets-OCR-s" print("Loading Nanonets OCR model...") model = AutoModelForImageTextToText.from_pretrained( model_path, torch_dtype="auto", device_map="auto", attn_implementation="flash_attention_2" ) model.eval() processor = AutoProcessor.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) print("Model loaded successfully!") # --- Helper Functions --- def process_tags(content: str) -> str: """Replaces special tags with HTML entities to prevent them from being rendered as HTML.""" content = content.replace("", "<img>") content = content.replace("", "</img>") content = content.replace("", "<watermark>") content = content.replace("", "</watermark>") content = content.replace("", "<page_number>") content = content.replace("", "</page_number>") content = content.replace("", "<signature>") content = content.replace("", "</signature>") return content def encode_image(image: Image) -> str: """Encodes an image to a base64 string.""" buffered = BytesIO() image.save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") return img_str @spaces.GPU def stream_request( messages: list[dict], model_name: str, max_tokens: int = 8000, temperature: float = 0.0, ): """ Stream text generation from the OCR model given messages with images and text. Args: messages: List of message dictionaries with role and content model_name: Name of the model (unused but kept for compatibility) max_tokens: Maximum number of tokens to generate temperature: Temperature for generation (unused, model runs deterministically) Yields: str: Generated text chunks """ # Extract the image and text from messages for message in messages: if message["role"] == "user": content = message["content"] image_data = None text_prompt = "" for item in content: if item["type"] == "image_url": # Decode base64 image image_url = item["image_url"]["url"] if image_url.startswith("data:image/jpeg;base64,"): image_base64 = image_url.split(",")[1] image_bytes = base64.b64decode(image_base64) image_data = Image.open(BytesIO(image_bytes)) elif item["type"] == "text": text_prompt = item["text"] if image_data is not None: # Format messages in the expected format for the model formatted_messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": [ {"type": "image", "image": image_data}, {"type": "text", "text": text_prompt}, ]}, ] # Apply chat template to format the input properly text = processor.apply_chat_template( formatted_messages, tokenize=False, add_generation_prompt=True ) # Process the formatted text and image inputs = processor( text=[text], images=[image_data], padding=True, return_tensors="pt" ) # Move inputs to the same device as the model inputs = {k: v.to(model.device) if hasattr(v, 'to') else v for k, v in inputs.items()} # Set up streaming streamer = TextIteratorStreamer( tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True ) generation_kwargs = { **inputs, "streamer": streamer, "max_new_tokens": max_tokens, "do_sample": False, # Deterministic generation "pad_token_id": tokenizer.eos_token_id, } # Start generation in a separate thread thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() # Yield generated tokens as they come for new_text in streamer: yield new_text thread.join() return # If no valid image/text pair found, return empty yield "" def convert_to_markdown_stream( images: Image, model_name, max_gen_tokens, with_img_desc: bool = False ): """ Generator function that yields streaming markdown conversion results Processes images one by one and concatenates results """ images = [images] # validate_file_paths(file_paths) # file_paths = convert_files_to_images(file_paths) # resize_images(file_paths, max_img_size) # Create system prompt for PDF to markdown conversion if with_img_desc: user_prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the tag; otherwise, add the image caption inside . Watermarks should be wrapped in brackets. Ex: OFFICIAL COPY. Page numbers should be wrapped in brackets. Ex: 14 or 9/22. Prefer using ☐ and ☑ for check boxes.""" else: user_prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Watermarks should be wrapped in brackets. Ex: OFFICIAL COPY. Page numbers should be wrapped in brackets. Ex: 14 or 9/22. Prefer using ☐ and ☑ for check boxes.""" # Accumulate results from all pages full_markdown_content = "" # Process each image individually for i, image in enumerate(images): # Build messages for this single image content = [ { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{encode_image(image)}" }, }, {"type": "text", "text": user_prompt}, ] messages = [{"role": "user", "content": content}] # Stream this individual page page_content = "" try: for chunk in stream_request( messages=messages, model_name=model_name, max_tokens=max_gen_tokens, ): page_content += chunk # Yield accumulated content from all pages processed so far + current page current_total = ( full_markdown_content + f"Page {i + 1} of {len(images)}\n" + page_content ) time.sleep(0.05) yield current_total # Process the completed page content and add it to the full content full_markdown_content += ( f"Page {i + 1} of {len(images)}\n" + page_content ) except Exception as e: return f"Error: {e}" def process_document(image, max_tokens, with_img_desc: bool = False): """ Process uploaded document (PDF or image) and convert to markdown. Args: file_path: Path to uploaded file max_tokens: Maximum tokens per page Returns: Generator yielding markdown content """ if image is None: return "Please upload a file first." try: # Handle PDF files # if file_path.name.lower().endswith('.pdf'): # # Convert PDF to images # with tempfile.TemporaryDirectory() as temp_dir: # # Copy uploaded file to temp directory # temp_pdf_path = os.path.join(temp_dir, "document.pdf") # import shutil # shutil.copy(file_path.name, temp_pdf_path) # # Convert PDF pages to images # images = convert_from_path(temp_pdf_path, dpi=150) # images = [image.convert("RGB") for image in images] # images = [image.resize((2048, 2048)) for image in images] # # Process each page # for result in convert_to_markdown_stream( # images, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc # ): # yield process_tags(result) # # Handle image files # else: # # Open image directly # image = Image.open(file_path.name).convert("RGB") # image = image.resize((2048, 2048)) image = Image.fromarray(image) image = image.resize((2048, 2048)) # Process single image for result in convert_to_markdown_stream( image, "nanonets/Nanonets-OCR-s", max_tokens, with_img_desc ): yield process_tags(result) except Exception as e: yield f"Error processing document: {str(e)}" # --- Gradio Interface --- with gr.Blocks(title="PDF to Markdown Converter", theme=gr.themes.Soft()) as demo: gr.HTML("""

📄 Nanonets-OCR-s: PDF & Image to Markdown Converter

Powered by Nanonets-OCR-s, A model for transforming documents into structured markdown with intelligent content recognition and semantic tagging.

📚 Hugging Face Model 📝 Release Blog 💻 GitHub Repository
""") with gr.Row(): with gr.Column(scale=1): file_input = gr.Image( label="Upload Image Document", height=200 ) max_tokens_slider = gr.Slider( minimum=1024, maximum=8192, value=4096, step=512, label="Max Tokens per Page", info="Maximum number of new tokens to generate for each page." ) with_img_desc_checkbox = gr.Checkbox( label="Include Image Description", value=False, info="If enabled, the model will include a description of the image in the output. If no image is present, use with_img_desc=False." ) extract_btn = gr.Button("Convert to Markdown", variant="primary", size="lg") with gr.Column(scale=2): output_text = gr.Markdown( label="Formatted Model Prediction", latex_delimiters=[{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}], line_breaks=True, show_copy_button=True, height=600, ) extract_btn.click( fn=process_document, inputs=[file_input, max_tokens_slider, with_img_desc_checkbox], concurrency_limit=4, outputs=output_text ) with gr.Accordion("About the Model (Nanonets-OCR-s)", open=False): gr.Markdown(""" ### Key Features - **LaTeX Equation Recognition**: Converts mathematical equations into properly formatted LaTeX. - **Intelligent Image Description**: Describes images within documents using structured `` tags. - **Signature & Watermark Detection**: Identifies and isolates signatures and watermarks within `` and `` tags. - **Smart Checkbox Handling**: Converts form checkboxes into standardized Unicode symbols (☐, ☑). - **Complex Table Extraction**: Accurately converts tables into HTML format. """) if __name__ == "__main__": demo.queue().launch(debug=True)