Update app.py

app.py

@@ -62,16 +62,24 @@
 # )
 
 # demo.launch()
-
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from PIL import Image
-import json
+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 smoldocling_readimage(image, prompt_text):
     messages = [
         {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]}
@@ -84,12 +92,8 @@ def smoldocling_readimage(image, prompt_text):
     result = processor.batch_decode(generated, skip_special_tokens=False)[0]
     clean_result = result.replace("<end_of_utterance>", "").strip()
 
-    # Try to parse as JSON
-    try:
-        json_result = json.loads(clean_result)
-        return json_result
-    except json.JSONDecodeError:
-        return {"error": "Output is not valid JSON", "raw_output": clean_result}
+    numbers = extract_numbers_from_docling(clean_result)
+    return numbers
 
 # Gradio UI
 demo = gr.Interface(
@@ -99,9 +103,8 @@ demo = gr.Interface(
         gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Convert to docling)", label="Prompt"),
     ],
     outputs=gr.JSON(),
-    title="SmolDocling Web App",
-    description="Upload a document image and convert it to structured docling format."
+    title="SmolDocling Web App - Extract Numbers",
+    description="Upload a document image and extract numeric values as a list."
 )
 
 demo.launch()
-