Update app.py

app.py

@@ -62,23 +62,30 @@
 # )
 
 # demo.launch()
+import re
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from PIL import Image
-import re
 
 # Load model & processor once at startup
 processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
 model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
 
-def extract_numbers_from_docling(docling_text):
-    # Remove tags except keep content between <fcel> and <nl>
-    # Use regex to find all numbers (integers or decimals)
-    numbers = re.findall(r"[-+]?\d*\.\d+|\d+", docling_text)
-    # Convert strings to floats or ints as appropriate
-    def convert_num(s):
-        return int(s) if s.isdigit() else float(s)
-    return [convert_num(num) for num in numbers]
+def extract_values(docling_text):
+    # Remove all <loc_*> tags
+    cleaned = re.sub(r"<loc_\d+>", "", docling_text)
+    # Split rows by <nl>
+    rows = cleaned.split("<nl>")
+    result = []
+    for row in rows:
+        if not row.strip():
+            continue
+        # Extract numbers inside <fcel> tags
+        values = re.findall(r"<fcel>(.*?)<fcel>", row)
+        # Convert to float list
+        float_values = [float(v) for v in values]
+        result.append(float_values)
+    return result
 
 def smoldocling_readimage(image, prompt_text):
     messages = [
@@ -89,11 +96,10 @@ def smoldocling_readimage(image, prompt_text):
     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]
-    clean_result = result.replace("<end_of_utterance>", "").strip()
-
-    numbers = extract_numbers_from_docling(clean_result)
-    return numbers
+    raw_result = processor.batch_decode(generated, skip_special_tokens=False)[0]
+    # Clean and extract numeric values
+    values_array = extract_values(raw_result)
+    return str(values_array)
 
 # Gradio UI
 demo = gr.Interface(
@@ -102,9 +108,9 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Upload Image"),
         gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Convert to docling)", label="Prompt"),
     ],
-    outputs=gr.JSON(),
-    title="SmolDocling Web App - Extract Numbers",
-    description="Upload a document image and extract numeric values as a list."
+    outputs="text",
+    title="SmolDocling Web App",
+    description="Upload a document image and convert it to structured docling format."
 )
 
 demo.launch()