Views
No views yet
ds4sd/SmolDocling-256M-preview using mlx-vlm version 0.1.18.mlx-vlm.Find Working MLX + Docling Example Code Below

1# Prerequisites:
2# pip install -U mlx-vlm
3# pip install docling_core
4
5import sys
6
7from pathlib import Path
8from PIL import Image
9
10from mlx_vlm import load, apply_chat_template, stream_generate
11from mlx_vlm.utils import load_image
12
13# Variables
14path_or_hf_repo="zboyles/SmolDocling-256M-preview-bf16"
15output_path=Path("output")
16output_path.mkdir(exist_ok=True)
17
18# Model Params
19eos="<end_of_utterance>"
20verbose=True
21kwargs={
22 "max_tokens": 8000,
23 "temperature": 0.0,
24}
25
26# Load images
27# Note: I manually downloaded the image
28# image_src = "https://upload.wikimedia.org/wikipedia/commons/7/76/GazettedeFrance.jpg"
29# image = load_image(image_src)
30image_src = "images/GazettedeFrance.jpg"
31image = Image.open(image_src).convert("RGB")
32
33# Initialize processor and model
34model, processor = load(
35 path_or_hf_repo=path_or_hf_repo,
36 trust_remote_code=True,
37)
38config = model.config
39
40
41# Create input messages - Docling Walkthrough Structure
42messages = [
43 {
44 "role": "user",
45 "content": [
46 {"type": "image"},
47 {"type": "text", "text": "Convert this page to docling."}
48 ]
49 },
50]
51prompt = apply_chat_template(processor, config, messages, add_generation_prompt=True)
52
53# # Alternatively, supported prompt creation method
54# messages = [{"role": "user", "content": "Convert this page to docling."}]
55# prompt = apply_chat_template(processor, config, messages, add_generation_prompt=True)
56
57
58text = ""
59last_response = None
60
61for response in stream_generate(
62 model=model,
63 processor=processor,
64 prompt=prompt,
65 image=image,
66 **kwargs
67):
68 if verbose:
69 print(response.text, end="", flush=True)
70 text += response.text
71 last_response = response
72 if eos in text:
73 text = text.split(eos)[0].strip()
74 break
75print()
76
77if verbose:
78 print("\n" + "=" * 10)
79 if len(text) == 0:
80 print("No text generated for this prompt")
81 sys.exit(0)
82 print(
83 f"Prompt: {last_response.prompt_tokens} tokens, "
84 f"{last_response.prompt_tps:.3f} tokens-per-sec"
85 )
86 print(
87 f"Generation: {last_response.generation_tokens} tokens, "
88 f"{last_response.generation_tps:.3f} tokens-per-sec"
89 )
90 print(f"Peak memory: {last_response.peak_memory:.3f} GB")
91
92# To convert to Docling Document, MD, HTML, etc.:
93docling_output_path = output_path / Path(image_src).with_suffix(".dt").name
94docling_output_path.write_text(text)
95doctags_doc = DocTagsDocument.from_doctags_and_image_pairs([text], [image])
96doc = DoclingDocument(name="Document")
97doc.load_from_doctags(doctags_doc)
98# export as any format
99# HTML
100doc.save_as_html(docling_output_path.with_suffix(".html"))
101# MD
102doc.save_as_markdown(docling_output_path.with_suffix(".md"))