Update app.py

app.py

@@ -1,28 +1,42 @@
+# app.py
+# Remember to add 'PyMuPDF', 'pdf2image', and 'torch' to your requirements.txt or install them.
+# For PDF processing, you might also need to install poppler:
+# On Debian/Ubuntu: sudo apt-get install poppler-utils
+# On macOS (using Homebrew): brew install poppler
+
 import gradio as gr
 from PIL import Image
-from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
-import torch
+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"
 
-# Load model once at startup
 print("Loading Nanonets OCR model...")
 model = AutoModelForImageTextToText.from_pretrained(
-    model_path, 
-    torch_dtype="auto", 
-    device_map="auto", 
+    model_path,
+    torch_dtype="auto",
+    device_map="auto",
     attn_implementation="flash_attention_2"
 )
 model.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(model_path)
 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>", "&lt;img&gt;")
     content = content.replace("</img>", "&lt;/img&gt;")
     content = content.replace("<watermark>", "&lt;watermark&gt;")
@@ -31,115 +45,229 @@ def process_tags(content: str) -> str:
     content = content.replace("</page_number>", "&lt;/page_number&gt;")
     content = content.replace("<signature>", "&lt;signature&gt;")
     content = content.replace("</signature>", "&lt;/signature&gt;")
-
     return content
 
-@spaces.GPU()
-def ocr_image_gradio_stream(image, max_tokens=4096):
-    """Process image through Nanonets OCR model with streaming output"""
-    if image is None:
-        yield "Please upload an image.", gr.update(interactive=True)
-        return
+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.
     
-    # Disable button at start of processing
-    yield "Processing image...", gr.update(interactive=False)
+    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)
     
-    try:
-        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: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
-        
-        # Convert PIL image if needed
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt},
-            ]},
-        ]
-        
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
-        inputs = inputs.to(model.device)
-        
-        # Set up streaming
-        streamer = TextIteratorStreamer(
-            tokenizer, 
-            timeout=60.0, 
-            skip_prompt=True, 
-            skip_special_tokens=True
-        )
-        
-        generation_kwargs = {
-            **inputs,
-            "max_new_tokens": max_tokens,
-            "do_sample": False,
-            "streamer": streamer
-        }
-        
-        # Start generation in a separate thread
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-        
-        # Stream the results
-        generated_text = ""
-        for new_text in streamer:
-            generated_text += new_text
-            processed_text = process_tags(generated_text)
-            yield process_tags(processed_text), gr.update(interactive=False)
-        
-        # Re-enable button when complete
-        yield process_tags(processed_text), gr.update(interactive=True)
+    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 = ""
             
-    except Exception as e:
-        error_msg = f"Error processing image: {str(e)}"
-        yield error_msg, gr.update(interactive=True)
-
-# Keep the original function for non-streaming use if needed
-@spaces.GPU()
-def ocr_image_gradio(image, max_tokens=4096):
-    """Process image through Nanonets OCR model for Gradio interface (non-streaming)"""
-    if image is None:
-        return "Please upload an image."
+            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
     
-    try:
-        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: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
-        
-        # Convert PIL image if needed
-        if not isinstance(image, Image.Image):
-            image = Image.fromarray(image)
-        
-        messages = [
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompt},
-            ]},
+    # 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 <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. 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: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. 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(file_path, 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 file_path 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)
         
-        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
-        inputs = inputs.to(model.device)
-        
-        with torch.no_grad():
-            output_ids = model.generate(**inputs, max_new_tokens=max_tokens, do_sample=False)
-            generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
+        # Handle image files
+        else:
+            # Open image directly
+            image = Image.open(file_path.name).convert("RGB")
+            image = image.resize((2048, 2048))
             
-        output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        return process_tags(output_text[0])
-        
+            # 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:
-        return f"Error processing image: {str(e)}"
+        yield f"Error processing document: {str(e)}"
 
-# Create Gradio interface
-with gr.Blocks(title="Nanonets OCR Demo") as demo:
-    # Replace simple markdown with styled HTML header that includes resources
+# --- Gradio Interface ---
+with gr.Blocks(title="PDF to Markdown Converter", theme=gr.themes.Soft()) as demo:
     gr.HTML("""
     <div class="title" style="text-align: center">
-        <h1>🔍 Nanonets OCR - Document Text Extraction</h1>
+        <h1>📄 Nanonets-OCR-s: PDF & Image to Markdown Converter</h1>
         <p style="font-size: 1.1em; color: #6b7280; margin-bottom: 0.6em;">
-            A model for transforming documents into structured markdown with intelligent content recognition and semantic tagging
+            Powered by <strong>Nanonets-OCR-s</strong>, A model for transforming documents into structured markdown with intelligent content recognition and semantic tagging.
         </p>
         <div style="display: flex; justify-content: center; gap: 20px; margin: 15px 0;">
             <a href="https://huggingface.co/nanonets/Nanonets-OCR-s" target="_blank" style="text-decoration: none; color: #2563eb; font-weight: 500;">
@@ -154,84 +282,54 @@ with gr.Blocks(title="Nanonets OCR Demo") as demo:
         </div>
     </div>
     """)
-    
+
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(
-                label="Upload Document Image",
-                type="pil",
-                height=400
+            file_input = gr.File(
+                label="Upload PDF or Image Document",
+                file_types=["pdf", "image"],
+                height=200
             )
             max_tokens_slider = gr.Slider(
                 minimum=1024,
                 maximum=8192,
                 value=4096,
                 step=512,
-                label="Max Tokens",
-                info="Maximum number of tokens to generate"
+                label="Max Tokens per Page",
+                info="Maximum number of new tokens to generate for each page."
             )
-            extract_btn = gr.Button("Extract Text", variant="primary", size="lg")
-            
+            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},
-                    {
-                        "left": "\\begin{align*}",
-                        "right": "\\end{align*}",
-                        "display": True,
-                    },
-                ],
+                label="Formatted Model Prediction",
+                latex_delimiters=[{"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}],
                 line_breaks=True,
                 show_copy_button=True,
+                height=600,
             )
-    
-    # Event handlers - Updated to use streaming
+
     extract_btn.click(
-        fn=ocr_image_gradio_stream,
-        inputs=[image_input, max_tokens_slider],
-        outputs=[output_text, extract_btn],
-        show_progress=True
+        fn=process_document,
+        inputs=[file_input, max_tokens_slider, with_img_desc_checkbox],
+        concurrency_limit=4,
+        outputs=output_text
     )
-    
-    # image_input.change(
-    #     fn=ocr_image_gradio_stream,
-    #     inputs=[image_input, max_tokens_slider],
-    #     outputs=[output_text, extract_btn],
-    #     show_progress=True
-    # )
-    
-    # Add model information section
-    with gr.Accordion("About Nanonets-OCR-s", open=False):
-        gr.Markdown("""
-## Nanonets-OCR-s
-
-Nanonets-OCR-s is a powerful, state-of-the-art image-to-markdown OCR model that goes far beyond traditional text extraction. 
-It transforms documents into structured markdown with intelligent content recognition and semantic tagging, making it ideal 
-for downstream processing by Large Language Models (LLMs).
-
-### Key Features
 
-- **LaTeX Equation Recognition**: Automatically converts mathematical equations and formulas into properly formatted LaTeX syntax. 
-    It distinguishes between inline `($...$)` and display `($$...$$)` equations.
-
-- **Intelligent Image Description**: Describes images within documents using structured `<img>` tags, making them digestible 
-    for LLM processing. It can describe various image types, including logos, charts, graphs and so on, detailing their content, 
-    style, and context.
-
-- **Signature Detection & Isolation**: Identifies and isolates signatures from other text, outputting them within a `<signature>` tag. 
-    This is crucial for processing legal and business documents.
-
-- **Watermark Extraction**: Detects and extracts watermark text from documents, placing it within a `<watermark>` tag.
-
-- **Smart Checkbox Handling**: Converts form checkboxes and radio buttons into standardized Unicode symbols (☐, ☑, ☒) 
-    for consistent and reliable processing.
-
-- **Complex Table Extraction**: Accurately extracts complex tables from documents and converts them into both markdown 
-    and HTML table formats.
-""")
+    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 `<img>` tags.
+        - **Signature & Watermark Detection**: Identifies and isolates signatures and watermarks within `<signature>` and `<watermark>` 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()
+    demo.queue().launch(debug=True)