fix: move to VLLM

backends/smoldocling.py

@@ -3,34 +3,28 @@
 # pip install docling_core
 # pip install transformers
 
-import torch
 from docling_core.types.doc import DoclingDocument
 from docling_core.types.doc.document import DocTagsDocument
 from img2table.document import PDF
 from PIL import Image
-from transformers import AutoModelForVision2Seq, AutoProcessor
+from vllm import LLM, SamplingParams
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Configuration
 MAX_PAGES = 1
+MODEL_PATH = "ds4sd/SmolDocling-256M-preview"
+PROMPT_TEXT = "Convert page to Docling."
 
 # Initialize processor and model
-processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
-model = AutoModelForVision2Seq.from_pretrained(
-    "ds4sd/SmolDocling-256M-preview",
-    torch_dtype=torch.bfloat16,
-    _attn_implementation="eager",
-).to(DEVICE)
+# Initialize LLM
+llm = LLM(
+    model=MODEL_PATH, limit_mm_per_prompt={"image": 1}, gpu_memory_utilization=0.4
+)
+chat_template = f"<|im_start|>User:<image>{PROMPT_TEXT}<end_of_utterance>\nAssistant:"
 
-# Create input messages
-messages = [
-    {
-        "role": "user",
-        "content": [
-            {"type": "image"},
-            {"type": "text", "text": "Convert this page to docling."},
-        ],
-    },
-]
+sampling_params = SamplingParams(
+    temperature=0.0,
+    max_tokens=4096,
+)
 
 
 def convert_smoldocling(path: str, file_name: str):
@@ -48,18 +42,10 @@ def convert_smoldocling(path: str, file_name: str):
             )
 
         # Prepare inputs
-        prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-        inputs = processor(text=prompt, images=[image], return_tensors="pt")
-        inputs = inputs.to(DEVICE)
+        llm_input = {"prompt": chat_template, "multi_modal_data": {"image": image}}
+        output = llm.generate([llm_input], sampling_params=sampling_params)[0]
 
-        # Generate outputs
-        generated_ids = model.generate(**inputs, max_new_tokens=4096)
-        prompt_length = inputs.input_ids.shape[1]
-        trimmed_generated_ids = generated_ids[:, prompt_length:]
-        doctags = processor.batch_decode(
-            trimmed_generated_ids,
-            skip_special_tokens=False,
-        )[0].lstrip()
+        doctags = output.outputs[0].text.lstrip()
 
         # Populate document
         doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([doctags], [image])

requirements.txt

@@ -11,3 +11,4 @@ docling_core
 opencv-contrib-python
 huggingface_hub
 spaces
+vllm