Update app.py

app.py

@@ -1,42 +1,38 @@
+# app.py
 import gradio as gr
-from PIL import Image
-from mlx_vlm import load, stream_generate
-from mlx_vlm.prompt_utils import apply_chat_template
-from mlx_vlm.utils import load_config
-from pdf2image import convert_from_path
+import pdfplumber
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-MODEL_PATH = "ibm-granite/granite-docling-258M-mlx"
+# Load the model
+model_name = "ibm-granite/granite-docling-258m-demo"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Load model once on startup
-model, processor = load(MODEL_PATH)
-config = load_config(MODEL_PATH)
+# Function to extract text from PDF
+def extract_text_from_pdf(pdf_file):
+    text = ""
+    with pdfplumber.open(pdf_file.name) as pdf:
+        for page in pdf.pages:
+            page_text = page.extract_text()
+            if page_text:
+                text += page_text + "\n"
+    return text
 
-def process_file(file):
-    # Convert PDF to images if needed
-    if file.name.endswith(".pdf"):
-        pages = convert_from_path(file.name)
-    else:
-        pages = [Image.open(file)]
-    
-    results = []
-    for page in pages:
-        prompt = "Convert this page to Markdown."
-        formatted_prompt = apply_chat_template(processor, config, prompt, num_images=1)
-        output = ""
-        for token in stream_generate(model, processor, formatted_prompt, [page], max_tokens=4096, verbose=False):
-            output += token.text
-            if "</doctag>" in token.text:
-                break
-        results.append(output)
-    
-    return "\n\n".join(results)
+# Function to generate JSON from text
+def pdf_to_json(pdf_file):
+    text = extract_text_from_pdf(pdf_file)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=2048)
+    outputs = model.generate(**inputs, max_new_tokens=1024)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result
 
-iface = gr.Interface(
-    fn=process_file,
-    inputs=gr.File(file_types=[".pdf", ".png", ".jpg", ".jpeg"]),
-    outputs=gr.Textbox(lines=30),
-    title="Docling PDF/Image to Markdown",
-    description="Upload a PDF or image. The model converts it to Markdown/structured format."
+# Gradio interface
+interface = gr.Interface(
+    fn=pdf_to_json,
+    inputs=gr.File(file_types=[".pdf"]),
+    outputs=gr.Textbox(label="Generated JSON"),
+    title="PDF to JSON using Granite DocLing",
+    description="Upload a PDF and get a JSON output using the ibm-granite/granite-docling-258m-demo model."
 )
 
-iface.launch()
+interface.launch()