mdipierro
/
SmolDocling-256M-Demo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix redundancies

This commit is contained in:
Maurizio Dipierro 2025-03-19 17:03:03 +01:00
parent 2ffed2f71e
commit 6023f099f8
1 changed files with 31 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
app.py
View File

@ -6,7 +6,6 @@ import re
import time import time
import torch import torch
import spaces import spaces
import re
import ast import ast
import html import html
import random import random
@ -18,18 +17,13 @@ from docling_core.types.doc.document import DocTagsDocument
def add_random_padding(image, min_percent=0.1, max_percent=0.10): def add_random_padding(image, min_percent=0.1, max_percent=0.10):
image = image.convert("RGB") image = image.convert("RGB")
width, height = image.size width, height = image.size
pad_w_percent = random.uniform(min_percent, max_percent) pad_w_percent = random.uniform(min_percent, max_percent)
pad_h_percent = random.uniform(min_percent, max_percent) pad_h_percent = random.uniform(min_percent, max_percent)
pad_w = int(width * pad_w_percent) pad_w = int(width * pad_w_percent)
pad_h = int(height * pad_h_percent) pad_h = int(height * pad_h_percent)
corner_pixel = image.getpixel((0, 0)) # Top-left corner corner_pixel = image.getpixel((0, 0)) # Top-left corner
padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel) padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel)
return padded_image return padded_image
def normalize_values(text, target_max=500): def normalize_values(text, target_max=500):
@ -46,16 +40,13 @@ def normalize_values(text, target_max=500):
normalized_text = re.sub(pattern, process_match, text) normalized_text = re.sub(pattern, process_match, text)
return normalized_text return normalized_text
processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview") processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
model = AutoModelForVision2Seq.from_pretrained("ds4sd/SmolDocling-256M-preview", model = AutoModelForVision2Seq.from_pretrained("ds4sd/SmolDocling-256M-preview",
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
#_attn_implementation="flash_attention_2" # _attn_implementation="flash_attention_2"
).to("cuda") ).to("cuda")
def model_inference( def model_inference(input_dict, history):
input_dict, history
):
text = input_dict["text"] text = input_dict["text"]
print(input_dict["files"]) print(input_dict["files"])
if len(input_dict["files"]) > 1: if len(input_dict["files"]) > 1:
@ -63,21 +54,18 @@ def model_inference(
images = [add_random_padding(load_image(image)) for image in input_dict["files"]] images = [add_random_padding(load_image(image)) for image in input_dict["files"]]
else: else:
images = [load_image(image) for image in input_dict["files"]] images = [load_image(image) for image in input_dict["files"]]
elif len(input_dict["files"]) == 1: elif len(input_dict["files"]) == 1:
if "OTSL" in text or "code" in text: if "OTSL" in text or "code" in text:
images = [add_random_padding(load_image(input_dict["files"][0]))] images = [add_random_padding(load_image(input_dict["files"][0]))]
else: else:
images = [load_image(input_dict["files"][0])] images = [load_image(input_dict["files"][0])]
else: else:
images = [] images = []
if text == "" and not images: if text == "" and not images:
gr.Error("Please input a query and optionally image(s).") gr.Error("Please input a query and optionally image(s).")
if text == "" and images: if text == "" and images:
gr.Error("Please input a text query along the image(s).") gr.Error("Please input a text query along with the image(s).")
if "OCR at text at" in text or "Identify element" in text or "formula" in text: if "OCR at text at" in text or "Identify element" in text or "formula" in text:
text = normalize_values(text, target_max=500) text = normalize_values(text, target_max=500)
@ -85,9 +73,7 @@ def model_inference(
resulting_messages = [ resulting_messages = [
{ {
"role": "user", "role": "user",
"content": [{"type": "image"} for _ in range(len(images))] + [ "content": [{"type": "image"} for _ in range(len(images))] + [{"type": "text", "text": text}]
{"type": "text", "text": text}
]
} }
] ]
prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True) prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
@ -116,12 +102,11 @@ def model_inference(
buffer += html.escape(new_text) buffer += html.escape(new_text)
yield buffer yield buffer
# After finishing the streamer loop:
cleaned_output = full_output.replace("<end_of_utterance>", "").strip() cleaned_output = full_output.replace("<end_of_utterance>", "").strip()
if cleaned_output: if cleaned_output:
yield cleaned_output yield cleaned_output
# Now, since cleaned_output exists, we can safely use it.
if any(tag in cleaned_output for tag in ["<doctag>", "<otsl>", "<code>", "<chart>", "<formula>"]): if any(tag in cleaned_output for tag in ["<doctag>", "<otsl>", "<code>", "<chart>", "<formula>"]):
doctag_output = cleaned_output doctag_output = cleaned_output
if "<chart>" in doctag_output: if "<chart>" in doctag_output:
@ -133,19 +118,8 @@ def model_inference(
doc.load_from_doctags(doctags_doc) doc.load_from_doctags(doctags_doc)
yield f"**MD Output:**\n\n{doc.export_to_markdown()}" yield f"**MD Output:**\n\n{doc.export_to_markdown()}"
examples = [
[{"text": "Convert this page to docling.", "files": ["example_images/2d0fbcc50e88065a040a537b717620e964fb4453314b71d83f3ed3425addcef6.png"]}],
if any(tag in doctag_output for tag in ["<doctag>", "<otsl>", "<code>", "<chart>", "<formula>"]):
doc = DoclingDocument(name="Document")
if "<chart>" in doctag_output:
doctag_output = doctag_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>")
doctag_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', doctag_output)
doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctag_output], images)
doc.load_from_doctags(doctags_doc)
yield f"**MD Output:**\n\n{doc.export_to_markdown()}"
examples=[[{"text": "Convert this page to docling.", "files": ["example_images/2d0fbcc50e88065a040a537b717620e964fb4453314b71d83f3ed3425addcef6.png"]}],
[{"text": "Convert this table to OTSL.", "files": ["example_images/image-2.jpg"]}], [{"text": "Convert this table to OTSL.", "files": ["example_images/image-2.jpg"]}],
[{"text": "Convert code to text.", "files": ["example_images/7666.jpg"]}], [{"text": "Convert code to text.", "files": ["example_images/7666.jpg"]}],
[{"text": "Convert formula to latex.", "files": ["example_images/2433.jpg"]}], [{"text": "Convert formula to latex.", "files": ["example_images/2433.jpg"]}],
@ -154,13 +128,17 @@ examples=[[{"text": "Convert this page to docling.", "files": ["example_images/2
[{"text": "Extract all section header elements on the page.", "files": ["example_images/paper_3.png"]}], [{"text": "Extract all section header elements on the page.", "files": ["example_images/paper_3.png"]}],
[{"text": "Identify element at location [123, 413, 1059, 1061]", "files": ["example_images/redhat.png"]}], [{"text": "Identify element at location [123, 413, 1059, 1061]", "files": ["example_images/redhat.png"]}],
[{"text": "Convert this page to docling.", "files": ["example_images/gazette_de_france.jpg"]}], [{"text": "Convert this page to docling.", "files": ["example_images/gazette_de_france.jpg"]}],
] ]
demo = gr.ChatInterface(fn=model_inference, title="SmolDocling-256M: Ultra-compact VLM for Document Conversion 💫", demo = gr.ChatInterface(
fn=model_inference,
title="SmolDocling-256M: Ultra-compact VLM for Document Conversion 💫",
description="Play with [ds4sd/SmolDocling-256M-preview](https://huggingface.co/ds4sd/SmolDocling-256M-preview) in this demo. To get started, upload an image and text or try one of the examples. This demo doesn't use history for the chat, so every chat you start is a new conversation.", description="Play with [ds4sd/SmolDocling-256M-preview](https://huggingface.co/ds4sd/SmolDocling-256M-preview) in this demo. To get started, upload an image and text or try one of the examples. This demo doesn't use history for the chat, so every chat you start is a new conversation.",
examples=examples, examples=examples,
textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"), stop_btn="Stop Generation", multimodal=True, textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple"),
stop_btn="Stop Generation",
multimodal=True,
cache_examples=False cache_examples=False
) )
demo.launch(debug=True,server_name="0.0.0.0", server_port=7860) demo.launch(debug=True, server_name="0.0.0.0", server_port=7860)