Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoProcessor, AutoModelForImageTextToText | |
| from PIL import Image | |
| # Load model & processor once at startup | |
| processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview") | |
| model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview") | |
| def smoldocling_readimage(image, prompt_text): | |
| """ | |
| Extract text and structured content from document images using SmolDocling model. | |
| This function processes document images (PDFs, scanned documents, screenshots, etc.) | |
| and converts them to structured text format based on the provided prompt. It uses | |
| the SmolDocling-256M-preview model for image-to-text conversion with chat-based | |
| prompting. | |
| Args: | |
| image (PIL.Image.Image): The input document image to process. Should be a PIL | |
| Image object containing a document, text, or any visual content that needs | |
| to be converted to text. | |
| prompt_text (str): The instruction or prompt text that guides the model's | |
| output format. Supported prompts include: | |
| Content Conversion: | |
| - "Convert this page to docling." - Full conversion to DocTags representation | |
| - "Convert chart to table." - Convert charts to table format | |
| - "Convert formula to LaTeX." - Convert mathematical formulas to LaTeX | |
| - "Convert code to text." - Convert code blocks to readable text | |
| - "Convert table to OTSL." - Convert tables to OTSL format (Lysak et al., 2023) | |
| OCR and Location-based Actions: | |
| - "OCR the text in a specific location: <loc_155><loc_233><loc_206><loc_237>" - Extract text from specific coordinates | |
| - "Identify element at: <loc_247><loc_482><loc_252><loc_486>" - Identify element type at coordinates | |
| - "Find all 'text' elements on the page, retrieve all section headers." - Extract section headers | |
| - "Detect footer elements on the page." - Identify footer content | |
| Returns: | |
| str: The extracted and formatted text content from the image, cleaned of | |
| special tokens and whitespace. The format depends on the prompt_text | |
| provided. | |
| Example: | |
| >>> from PIL import Image | |
| >>> img = Image.open("document.pdf") | |
| >>> result = smoldocling_readimage(img, "Convert to docling") | |
| >>> print(result) # Returns structured document content | |
| Note: | |
| - The function is optimized for document images but can handle any image | |
| containing text | |
| - Processing time depends on image size and complexity | |
| - Maximum output length is limited to 1024 new tokens | |
| """ | |
| messages = [ | |
| {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]} | |
| ] | |
| prompt = processor.apply_chat_template(messages, add_generation_prompt=True) | |
| inputs = processor(text=prompt, images=[image], return_tensors="pt") | |
| outputs = model.generate(**inputs, max_new_tokens=1024) | |
| prompt_length = inputs.input_ids.shape[1] | |
| generated = outputs[:, prompt_length:] | |
| result = processor.batch_decode(generated, skip_special_tokens=False)[0] | |
| return result.replace("<end_of_utterance>", "").strip() | |
| # Gradio UI | |
| demo = gr.Interface( | |
| fn=smoldocling_readimage, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Image"), | |
| gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Convert to docling)", label="Prompt"), | |
| ], | |
| outputs="text", | |
| title="SmolDocling Web App", | |
| description="Upload a document image and convert it to structured docling format." | |
| ) | |
| demo.launch(mcp_server=True) |