beaupi commited on
Commit
d21ec71
·
verified ·
1 Parent(s): 72ea311

Upload dots.ocr-oQ8 via oMLX

Browse files
README.md ADDED
@@ -0,0 +1,1315 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ library_name: dots_ocr
4
+ pipeline_tag: image-text-to-text
5
+ tags:
6
+ - image-to-text
7
+ - ocr
8
+ - document-parse
9
+ - layout
10
+ - table
11
+ - formula
12
+ - transformers
13
+ - custom_code
14
+ language:
15
+ - en
16
+ - zh
17
+ - multilingual
18
+ ---
19
+
20
+ <div align="center">
21
+
22
+ <p align="center">
23
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/logo.png" width="300"/>
24
+ <p>
25
+
26
+ <h1 align="center">
27
+ dots.ocr: Multilingual Document Layout Parsing in a Single Vision-Language Model
28
+ </h1>
29
+
30
+ [![Blog](https://img.shields.io/badge/Blog-View_on_GitHub-333.svg?logo=github)](https://github.com/rednote-hilab/dots.ocr/blob/master/assets/blog.md)
31
+ [![HuggingFace](https://img.shields.io/badge/HuggingFace%20Weights-black.svg?logo=HuggingFace)](https://huggingface.co/rednote-hilab/dots.ocr)
32
+
33
+
34
+ <div align="center">
35
+ <a href="https://dotsocr.xiaohongshu.com" target="_blank" rel="noopener noreferrer"><strong>🖥️ Live Demo</strong></a> |
36
+ <a href="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/wechat.jpg" target="_blank" rel="noopener noreferrer"><strong>💬 WeChat</strong></a> |
37
+ <a href="https://www.xiaohongshu.com/user/profile/683ffe42000000001d021a4c" target="_blank" rel="noopener noreferrer"><strong>📕 rednote</strong></a>
38
+ </div>
39
+
40
+ </div>
41
+
42
+
43
+
44
+ ## Introduction
45
+
46
+ **dots.ocr** is a powerful, multilingual document parser that unifies layout detection and content recognition within a single vision-language model while maintaining good reading order. Despite its compact 1.7B-parameter LLM foundation, it achieves state-of-the-art(SOTA) performance.
47
+
48
+ 1. **Powerful Performance:** **dots.ocr** achieves SOTA performance for text, tables, and reading order on [OmniDocBench](https://github.com/opendatalab/OmniDocBench), while delivering formula recognition results comparable to much larger models like Doubao-1.5 and gemini2.5-pro.
49
+ 2. **Multilingual Support:** **dots.ocr** demonstrates robust parsing capabilities for low-resource languages, achieving decisive advantages across both layout detection and content recognition on our in-house multilingual documents benchmark.
50
+ 3. **Unified and Simple Architecture:** By leveraging a single vision-language model, **dots.ocr** offers a significantly more streamlined architecture than conventional methods that rely on complex, multi-model pipelines. Switching between tasks is accomplished simply by altering the input prompt, proving that a VLM can achieve competitive detection results compared to traditional detection models like DocLayout-YOLO.
51
+ 4. **Efficient and Fast Performance:** Built upon a compact 1.7B LLM, **dots.ocr** provides faster inference speeds than many other high-performing models based on larger foundations.
52
+
53
+
54
+ ## Usage with transformers
55
+
56
+ ```py
57
+ import torch
58
+ from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
59
+ from qwen_vl_utils import process_vision_info
60
+ from dots_ocr.utils import dict_promptmode_to_prompt
61
+
62
+ model_path = "./weights/DotsOCR"
63
+ model = AutoModelForCausalLM.from_pretrained(
64
+ model_path,
65
+ attn_implementation="flash_attention_2",
66
+ torch_dtype=torch.bfloat16,
67
+ device_map="auto",
68
+ trust_remote_code=True
69
+ )
70
+ processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
71
+
72
+ image_path = "demo/demo_image1.jpg"
73
+ prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
74
+
75
+ 1. Bbox format: [x1, y1, x2, y2]
76
+
77
+ 2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
78
+
79
+ 3. Text Extraction & Formatting Rules:
80
+ - Picture: For the 'Picture' category, the text field should be omitted.
81
+ - Formula: Format its text as LaTeX.
82
+ - Table: Format its text as HTML.
83
+ - All Others (Text, Title, etc.): Format their text as Markdown.
84
+
85
+ 4. Constraints:
86
+ - The output text must be the original text from the image, with no translation.
87
+ - All layout elements must be sorted according to human reading order.
88
+
89
+ 5. Final Output: The entire output must be a single JSON object.
90
+ """
91
+
92
+ messages = [
93
+ {
94
+ "role": "user",
95
+ "content": [
96
+ {
97
+ "type": "image",
98
+ "image": image_path
99
+ },
100
+ {"type": "text", "text": prompt}
101
+ ]
102
+ }
103
+ ]
104
+
105
+ # Preparation for inference
106
+ text = processor.apply_chat_template(
107
+ messages,
108
+ tokenize=False,
109
+ add_generation_prompt=True
110
+ )
111
+ image_inputs, video_inputs = process_vision_info(messages)
112
+ inputs = processor(
113
+ text=[text],
114
+ images=image_inputs,
115
+ videos=video_inputs,
116
+ padding=True,
117
+ return_tensors="pt",
118
+ )
119
+
120
+ inputs = inputs.to("cuda")
121
+
122
+ # Inference: Generation of the output
123
+ generated_ids = model.generate(**inputs, max_new_tokens=24000)
124
+ generated_ids_trimmed = [
125
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
126
+ ]
127
+ output_text = processor.batch_decode(
128
+ generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
129
+ )
130
+ print(output_text)
131
+ ```
132
+
133
+ ### Performance Comparison: dots.ocr vs. Competing Models
134
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/chart.png" border="0" />
135
+
136
+ > **Notes:**
137
+ > - The EN, ZH metrics are the end2end evaluation results of [OmniDocBench](https://github.com/opendatalab/OmniDocBench), and Multilingual metric is the end2end evaluation results of dots.ocr-bench.
138
+
139
+
140
+ ## News
141
+ * ```2025.07.30 ``` 🚀 We release [dots.ocr](https://github.com/rednote-hilab/dots.ocr), — a multilingual documents parsing model based on 1.7b llm, with SOTA performance.
142
+
143
+
144
+
145
+ ## Benchmark Results
146
+
147
+ ### 1. OmniDocBench
148
+
149
+ #### The end-to-end evaluation results of different tasks.
150
+
151
+ <table>
152
+ <thead>
153
+ <tr>
154
+ <th rowspan="2"><strong>Model<br>Type</strong></th>
155
+ <th rowspan="2"><strong>Methods</strong></th>
156
+ <th colspan="2"><strong>Overall<sup>Edit</sup>↓</strong></th>
157
+ <th colspan="2"><strong>Text<sup>Edit</sup>↓</strong></th>
158
+ <th colspan="2"><strong>Formula<sup>Edit</sup>↓</strong></th>
159
+ <th colspan="2"><strong>Table<sup>TEDS</sup>↑</strong></th>
160
+ <th colspan="2"><strong>Table<sup>Edit</sup>↓</strong></th>
161
+ <th colspan="2"><strong>Read Order<sup>Edit</sup>↓</strong></th>
162
+ </tr>
163
+ <tr>
164
+ <th><em>EN</em></th>
165
+ <th><em>ZH</em></th>
166
+ <th><em>EN</em></th>
167
+ <th><em>ZH</em></th>
168
+ <th><em>EN</em></th>
169
+ <th><em>ZH</em></th>
170
+ <th><em>EN</em></th>
171
+ <th><em>ZH</em></th>
172
+ <th><em>EN</em></th>
173
+ <th><em>ZH</em></th>
174
+ <th><em>EN</em></th>
175
+ <th><em>ZH</em></th>
176
+ </tr>
177
+ </thead>
178
+ <tbody>
179
+ <tr>
180
+ <td rowspan="8"><strong>Pipeline<br>Tools</strong></td>
181
+ <td>MinerU</td>
182
+ <td>0.150</td>
183
+ <td>0.357</td>
184
+ <td>0.061</td>
185
+ <td>0.215</td>
186
+ <td>0.278</td>
187
+ <td>0.577</td>
188
+ <td>78.6</td>
189
+ <td>62.1</td>
190
+ <td>0.180</td>
191
+ <td>0.344</td>
192
+ <td>0.079</td>
193
+ <td>0.292</td>
194
+ </tr>
195
+ <tr>
196
+ <td>Marker</td>
197
+ <td>0.336</td>
198
+ <td>0.556</td>
199
+ <td>0.080</td>
200
+ <td>0.315</td>
201
+ <td>0.530</td>
202
+ <td>0.883</td>
203
+ <td>67.6</td>
204
+ <td>49.2</td>
205
+ <td>0.619</td>
206
+ <td>0.685</td>
207
+ <td>0.114</td>
208
+ <td>0.340</td>
209
+ </tr>
210
+ <tr>
211
+ <td>Mathpix</td>
212
+ <td>0.191</td>
213
+ <td>0.365</td>
214
+ <td>0.105</td>
215
+ <td>0.384</td>
216
+ <td>0.306</td>
217
+ <td>0.454</td>
218
+ <td>77.0</td>
219
+ <td>67.1</td>
220
+ <td>0.243</td>
221
+ <td>0.320</td>
222
+ <td>0.108</td>
223
+ <td>0.304</td>
224
+ </tr>
225
+ <tr>
226
+ <td>Docling</td>
227
+ <td>0.589</td>
228
+ <td>0.909</td>
229
+ <td>0.416</td>
230
+ <td>0.987</td>
231
+ <td>0.999</td>
232
+ <td>1</td>
233
+ <td>61.3</td>
234
+ <td>25.0</td>
235
+ <td>0.627</td>
236
+ <td>0.810</td>
237
+ <td>0.313</td>
238
+ <td>0.837</td>
239
+ </tr>
240
+ <tr>
241
+ <td>Pix2Text</td>
242
+ <td>0.320</td>
243
+ <td>0.528</td>
244
+ <td>0.138</td>
245
+ <td>0.356</td>
246
+ <td>0.276</td>
247
+ <td>0.611</td>
248
+ <td>73.6</td>
249
+ <td>66.2</td>
250
+ <td>0.584</td>
251
+ <td>0.645</td>
252
+ <td>0.281</td>
253
+ <td>0.499</td>
254
+ </tr>
255
+ <tr>
256
+ <td>Unstructured</td>
257
+ <td>0.586</td>
258
+ <td>0.716</td>
259
+ <td>0.198</td>
260
+ <td>0.481</td>
261
+ <td>0.999</td>
262
+ <td>1</td>
263
+ <td>0</td>
264
+ <td>0.06</td>
265
+ <td>1</td>
266
+ <td>0.998</td>
267
+ <td>0.145</td>
268
+ <td>0.387</td>
269
+ </tr>
270
+ <tr>
271
+ <td>OpenParse</td>
272
+ <td>0.646</td>
273
+ <td>0.814</td>
274
+ <td>0.681</td>
275
+ <td>0.974</td>
276
+ <td>0.996</td>
277
+ <td>1</td>
278
+ <td>64.8</td>
279
+ <td>27.5</td>
280
+ <td>0.284</td>
281
+ <td>0.639</td>
282
+ <td>0.595</td>
283
+ <td>0.641</td>
284
+ </tr>
285
+ <tr>
286
+ <td>PPStruct-V3</td>
287
+ <td>0.145</td>
288
+ <td>0.206</td>
289
+ <td>0.058</td>
290
+ <td>0.088</td>
291
+ <td>0.295</td>
292
+ <td>0.535</td>
293
+ <td>-</td>
294
+ <td>-</td>
295
+ <td>0.159</td>
296
+ <td>0.109</td>
297
+ <td>0.069</td>
298
+ <td>0.091</td>
299
+ </tr>
300
+ <tr>
301
+ <td rowspan="9"><strong>Expert<br>VLMs</strong></td>
302
+ <td>GOT-OCR</td>
303
+ <td>0.287</td>
304
+ <td>0.411</td>
305
+ <td>0.189</td>
306
+ <td>0.315</td>
307
+ <td>0.360</td>
308
+ <td>0.528</td>
309
+ <td>53.2</td>
310
+ <td>47.2</td>
311
+ <td>0.459</td>
312
+ <td>0.520</td>
313
+ <td>0.141</td>
314
+ <td>0.280</td>
315
+ </tr>
316
+ <tr>
317
+ <td>Nougat</td>
318
+ <td>0.452</td>
319
+ <td>0.973</td>
320
+ <td>0.365</td>
321
+ <td>0.998</td>
322
+ <td>0.488</td>
323
+ <td>0.941</td>
324
+ <td>39.9</td>
325
+ <td>0</td>
326
+ <td>0.572</td>
327
+ <td>1.000</td>
328
+ <td>0.382</td>
329
+ <td>0.954</td>
330
+ </tr>
331
+ <tr>
332
+ <td>Mistral OCR</td>
333
+ <td>0.268</td>
334
+ <td>0.439</td>
335
+ <td>0.072</td>
336
+ <td>0.325</td>
337
+ <td>0.318</td>
338
+ <td>0.495</td>
339
+ <td>75.8</td>
340
+ <td>63.6</td>
341
+ <td>0.600</td>
342
+ <td>0.650</td>
343
+ <td>0.083</td>
344
+ <td>0.284</td>
345
+ </tr>
346
+ <tr>
347
+ <td>OLMOCR-sglang</td>
348
+ <td>0.326</td>
349
+ <td>0.469</td>
350
+ <td>0.097</td>
351
+ <td>0.293</td>
352
+ <td>0.455</td>
353
+ <td>0.655</td>
354
+ <td>68.1</td>
355
+ <td>61.3</td>
356
+ <td>0.608</td>
357
+ <td>0.652</td>
358
+ <td>0.145</td>
359
+ <td>0.277</td>
360
+ </tr>
361
+ <tr>
362
+ <td>SmolDocling-256M</td>
363
+ <td>0.493</td>
364
+ <td>0.816</td>
365
+ <td>0.262</td>
366
+ <td>0.838</td>
367
+ <td>0.753</td>
368
+ <td>0.997</td>
369
+ <td>44.9</td>
370
+ <td>16.5</td>
371
+ <td>0.729</td>
372
+ <td>0.907</td>
373
+ <td>0.227</td>
374
+ <td>0.522</td>
375
+ </tr>
376
+ <tr>
377
+ <td>Dolphin</td>
378
+ <td>0.206</td>
379
+ <td>0.306</td>
380
+ <td>0.107</td>
381
+ <td>0.197</td>
382
+ <td>0.447</td>
383
+ <td>0.580</td>
384
+ <td>77.3</td>
385
+ <td>67.2</td>
386
+ <td>0.180</td>
387
+ <td>0.285</td>
388
+ <td>0.091</td>
389
+ <td>0.162</td>
390
+ </tr>
391
+ <tr>
392
+ <td>MinerU 2</td>
393
+ <td>0.139</td>
394
+ <td>0.240</td>
395
+ <td>0.047</td>
396
+ <td>0.109</td>
397
+ <td>0.297</td>
398
+ <td>0.536</td>
399
+ <td>82.5</td>
400
+ <td>79.0</td>
401
+ <td>0.141</td>
402
+ <td>0.195</td>
403
+ <td>0.069<</td>
404
+ <td>0.118</td>
405
+ </tr>
406
+ <tr>
407
+ <td>OCRFlux</td>
408
+ <td>0.195</td>
409
+ <td>0.281</td>
410
+ <td>0.064</td>
411
+ <td>0.183</td>
412
+ <td>0.379</td>
413
+ <td>0.613</td>
414
+ <td>71.6</td>
415
+ <td>81.3</td>
416
+ <td>0.253</td>
417
+ <td>0.139</td>
418
+ <td>0.086</td>
419
+ <td>0.187</td>
420
+ </tr>
421
+ <tr>
422
+ <td>MonkeyOCR-pro-3B</td>
423
+ <td>0.138</td>
424
+ <td>0.206</td>
425
+ <td>0.067</td>
426
+ <td>0.107</td>
427
+ <td><strong>0.246</strong></td>
428
+ <td>0.421</td>
429
+ <td>81.5</td>
430
+ <td>87.5</td>
431
+ <td>0.139</td>
432
+ <td>0.111</td>
433
+ <td>0.100</td>
434
+ <td>0.185</td>
435
+ </tr>
436
+ <tr>
437
+
438
+ <td rowspan="5"><strong>General<br>VLMs</strong></td>
439
+ <td>GPT4o</td>
440
+ <td>0.233</td>
441
+ <td>0.399</td>
442
+ <td>0.144</td>
443
+ <td>0.409</td>
444
+ <td>0.425</td>
445
+ <td>0.606</td>
446
+ <td>72.0</td>
447
+ <td>62.9</td>
448
+ <td>0.234</td>
449
+ <td>0.329</td>
450
+ <td>0.128</td>
451
+ <td>0.251</td>
452
+ </tr>
453
+ <tr>
454
+ <td>Qwen2-VL-72B</td>
455
+ <td>0.252</td>
456
+ <td>0.327</td>
457
+ <td>0.096</td>
458
+ <td>0.218</td>
459
+ <td>0.404</td>
460
+ <td>0.487</td>
461
+ <td>76.8</td>
462
+ <td>76.4</td>
463
+ <td>0.387</td>
464
+ <td>0.408</td>
465
+ <td>0.119</td>
466
+ <td>0.193</td>
467
+ </tr>
468
+ <tr>
469
+ <td>Qwen2.5-VL-72B</td>
470
+ <td>0.214</td>
471
+ <td>0.261</td>
472
+ <td>0.092</td>
473
+ <td>0.18</td>
474
+ <td>0.315</td>
475
+ <td>0.434</td>
476
+ <td>82.9</td>
477
+ <td>83.9</td>
478
+ <td>0.341</td>
479
+ <td>0.262</td>
480
+ <td>0.106</td>
481
+ <td>0.168</td>
482
+ </tr>
483
+ <tr>
484
+ <td>Gemini2.5-Pro</td>
485
+ <td>0.148</td>
486
+ <td>0.212</td>
487
+ <td>0.055</td>
488
+ <td>0.168</td>
489
+ <td>0.356</td>
490
+ <td>0.439</td>
491
+ <td>85.8</td>
492
+ <td>86.4</td>
493
+ <td>0.13</td>
494
+ <td>0.119</td>
495
+ <td>0.049</td>
496
+ <td>0.121</td>
497
+ </tr>
498
+ <tr>
499
+ <td>doubao-1-5-thinking-vision-pro-250428</td>
500
+ <td>0.140</td>
501
+ <td>0.162</td>
502
+ <td>0.043</td>
503
+ <td>0.085</td>
504
+ <td>0.295</td>
505
+ <td><strong>0.384</strong></td>
506
+ <td>83.3</td>
507
+ <td><strong>89.3</strong></td>
508
+ <td>0.165</td>
509
+ <td><strong>0.085</strong></td>
510
+ <td>0.058</td>
511
+ <td>0.094</td>
512
+ </tr>
513
+ <tr>
514
+ <td rowspan="1"><strong>Expert VLMs</strong></td>
515
+ <td><strong>dots.ocr</strong></td>
516
+ <td><strong>0.125</strong></td>
517
+ <td><strong>0.160</strong></td>
518
+ <td><strong>0.032</strong></td>
519
+ <td><strong>0.066</strong></td>
520
+ <td>0.329</td>
521
+ <td>0.416</td>
522
+ <td><strong>88.6</strong></td>
523
+ <td>89.0</td>
524
+ <td><strong>0.099</strong></td>
525
+ <td>0.092</td>
526
+ <td><strong>0.040</strong></td>
527
+ <td><strong>0.067</strong></td>
528
+ </tr>
529
+ <tr>
530
+ </tbody>
531
+ </table>
532
+
533
+
534
+ #### The end-to-end text recognition performance across 9 PDF page types.
535
+
536
+ <table>
537
+ <thead>
538
+ <tr>
539
+ <th><strong>Model<br>Type</strong></th>
540
+ <th><strong>Models</strong></th>
541
+ <th><strong>Book</strong></th>
542
+ <th><strong>Slides</strong></th>
543
+ <th><strong>Financial<br>Report</strong></th>
544
+ <th><strong>Textbook</strong></th>
545
+ <th><strong>Exam<br>Paper</strong></th>
546
+ <th><strong>Magazine</strong></th>
547
+ <th><strong>Academic<br>Papers</strong></th>
548
+ <th><strong>Notes</strong></th>
549
+ <th><strong>Newspaper</strong></th>
550
+ <th><strong>Overall</strong></th>
551
+ </tr>
552
+ </thead>
553
+ <tbody>
554
+ <tr>
555
+ <td rowspan="3"><strong>Pipeline<br>Tools</strong></td>
556
+ <td>MinerU</td>
557
+ <td>0.055</td>
558
+ <td>0.124</td>
559
+ <td><u>0.033</u></td>
560
+ <td>0.102</td>
561
+ <td>0.159</td>
562
+ <td><strong>0.072</strong></td>
563
+ <td><u>0.025</u></td>
564
+ <td>0.984</td>
565
+ <td>0.171</td>
566
+ <td>0.206</td>
567
+ </tr>
568
+ <tr>
569
+ <td>Marker</td>
570
+ <td>0.074</td>
571
+ <td>0.340</td>
572
+ <td>0.089</td>
573
+ <td>0.319</td>
574
+ <td>0.452</td>
575
+ <td>0.153</td>
576
+ <td>0.059</td>
577
+ <td>0.651</td>
578
+ <td>0.192</td>
579
+ <td>0.274</td>
580
+ </tr>
581
+ <tr>
582
+ <td>Mathpix</td>
583
+ <td>0.131</td>
584
+ <td>0.220</td>
585
+ <td>0.202</td>
586
+ <td>0.216</td>
587
+ <td>0.278</td>
588
+ <td>0.147</td>
589
+ <td>0.091</td>
590
+ <td>0.634</td>
591
+ <td>0.690</td>
592
+ <td>0.300</td>
593
+ </tr>
594
+ <tr>
595
+ <td rowspan="5"><strong>Expert<br>VLMs</strong></td>
596
+ <td>GOT-OCR</td>
597
+ <td>0.111</td>
598
+ <td>0.222</td>
599
+ <td>0.067</td>
600
+ <td>0.132</td>
601
+ <td>0.204</td>
602
+ <td>0.198</td>
603
+ <td>0.179</td>
604
+ <td>0.388</td>
605
+ <td>0.771</td>
606
+ <td>0.267</td>
607
+ </tr>
608
+ <tr>
609
+ <td>Nougat</td>
610
+ <td>0.734</td>
611
+ <td>0.958</td>
612
+ <td>1.000</td>
613
+ <td>0.820</td>
614
+ <td>0.930</td>
615
+ <td>0.830</td>
616
+ <td>0.214</td>
617
+ <td>0.991</td>
618
+ <td>0.871</td>
619
+ <td>0.806</td>
620
+ </tr>
621
+ <tr>
622
+ <td>Dolphin</td>
623
+ <td>0.091</td>
624
+ <td>0.131</td>
625
+ <td>0.057</td>
626
+ <td>0.146</td>
627
+ <td>0.231</td>
628
+ <td>0.121</td>
629
+ <td>0.074</td>
630
+ <td>0.363</td>
631
+ <td>0.307</td>
632
+ <td>0.177</td>
633
+ </tr>
634
+ <tr>
635
+ <td>OCRFlux</td>
636
+ <td>0.068</td>
637
+ <td>0.125</td>
638
+ <td>0.092</td>
639
+ <td>0.102</td>
640
+ <td>0.119</td>
641
+ <td>0.083</td>
642
+ <td>0.047</td>
643
+ <td>0.223</td>
644
+ <td>0.536</td>
645
+ <td>0.149</td>
646
+ </tr>
647
+ <tr>
648
+ <td>MonkeyOCR-pro-3B</td>
649
+ <td>0.084</td>
650
+ <td>0.129</td>
651
+ <td>0.060</td>
652
+ <td>0.090</td>
653
+ <td>0.107</td>
654
+ <td>0.073</td>
655
+ <td>0.050</td>
656
+ <td>0.171</td>
657
+ <td>0.107</td>
658
+ <td>0.100</td>
659
+ </tr>
660
+ <tr>
661
+ <td rowspan="4"><strong>General<br>VLMs</strong></td>
662
+ <td>GPT4o</td>
663
+ <td>0.157</td>
664
+ <td>0.163</td>
665
+ <td>0.348</td>
666
+ <td>0.187</td>
667
+ <td>0.281</td>
668
+ <td>0.173</td>
669
+ <td>0.146</td>
670
+ <td>0.607</td>
671
+ <td>0.751</td>
672
+ <td>0.316</td>
673
+ </tr>
674
+ <tr>
675
+ <td>Qwen2.5-VL-7B</td>
676
+ <td>0.148</td>
677
+ <td>0.053</td>
678
+ <td>0.111</td>
679
+ <td>0.137</td>
680
+ <td>0.189</td>
681
+ <td>0.117</td>
682
+ <td>0.134</td>
683
+ <td>0.204</td>
684
+ <td>0.706</td>
685
+ <td>0.205</td>
686
+ </tr>
687
+ <tr>
688
+ <td>InternVL3-8B</td>
689
+ <td>0.163</td>
690
+ <td>0.056</td>
691
+ <td>0.107</td>
692
+ <td>0.109</td>
693
+ <td>0.129</td>
694
+ <td>0.100</td>
695
+ <td>0.159</td>
696
+ <td>0.150</td>
697
+ <td>0.681</td>
698
+ <td>0.188</td>
699
+ </tr>
700
+ <tr>
701
+ <td>doubao-1-5-thinking-vision-pro-250428</td>
702
+ <td>0.048</td>
703
+ <td>0.048</td>
704
+ <td>0.024</td>
705
+ <td><strong>0.062</strong></td>
706
+ <td>0.085</td>
707
+ <td>0.051</td>
708
+ <td>0.039</td>
709
+ <td><strong>0.096</strong></td>
710
+ <td>0.181</td>
711
+ <td>0.073</td>
712
+ </tr>
713
+ <tr>
714
+ <td rowspan="1"><strong>Expert VLMs</strong></td>
715
+ <td><strong>dots.ocr</strong></td>
716
+ <td><strong>0.031</strong></td>
717
+ <td><strong>0.047</strong></td>
718
+ <td><strong>0.011</strong></td>
719
+ <td>0.082</td>
720
+ <td><strong>0.079</strong></td>
721
+ <td><strong>0.028</strong></td>
722
+ <td><strong>0.029</strong></td>
723
+ <td>0.109</td>
724
+ <td><strong>0.056</strong></td>
725
+ <td><strong>0.055</strong></td>
726
+ </tr>
727
+
728
+ </tbody>
729
+ </table>
730
+
731
+ > **Notes:**
732
+ > - The metrics are from [MonkeyOCR](https://github.com/Yuliang-Liu/MonkeyOCR), [OmniDocBench](https://github.com/opendatalab/OmniDocBench), and our own internal evaluations.
733
+ > - We delete the Page-header and Page-footer cells in the result markdown.
734
+ > - We use tikz_preprocess pipeline to upsample the images to dpi 200.
735
+
736
+
737
+ ### 2. **dots.ocr-bench**
738
+
739
+ This is an inhouse benchmark which contain 1493 pdf images with 100 languages.
740
+
741
+ #### The end-to-end evaluation results of different tasks.
742
+
743
+ <table>
744
+ <thead>
745
+ <tr>
746
+ <th rowspan="1"><strong>Methods</strong></th>
747
+ <th colspan="1"><strong>Overall<sup>Edit</sup>↓</strong></th>
748
+ <th colspan="1"><strong>Text<sup>Edit</sup>↓</strong></th>
749
+ <th colspan="1"><strong>Formula<sup>Edit</sup>↓</strong></th>
750
+ <th colspan="1"><strong>Table<sup>TEDS</sup>↑</strong></th>
751
+ <th colspan="1"><strong>Table<sup>Edit</sup>↓</strong></th>
752
+ <th colspan="1"><strong>Read Order<sup>Edit</sup>↓</strong></th>
753
+ </tr>
754
+ </thead>
755
+ <tbody>
756
+ <td>MonkeyOCR-3B</td>
757
+ <td>0.483</td>
758
+ <td>0.445</td>
759
+ <td>0.627</td>
760
+ <td>50.93</td>
761
+ <td>0.452</td>
762
+ <td>0.409</td>
763
+ </tr>
764
+ <tr>
765
+ <td>doubao-1-5-thinking-vision-pro-250428</td>
766
+ <td>0.291</td>
767
+ <td>0.226</td>
768
+ <td>0.440</td>
769
+ <td>71.2</td>
770
+ <td>0.260</td>
771
+ <td>0.238</td>
772
+ </tr>
773
+ <tr>
774
+ <td>doubao-1-6</td>
775
+ <td>0.299</td>
776
+ <td>0.270</td>
777
+ <td>0.417</td>
778
+ <td>71.0</td>
779
+ <td>0.258</td>
780
+ <td>0.253</td>
781
+ </tr>
782
+ <tr>
783
+ <td>Gemini2.5-Pro</td>
784
+ <td>0.251</td>
785
+ <td>0.163</td>
786
+ <td>0.402</td>
787
+ <td>77.1</td>
788
+ <td>0.236</td>
789
+ <td>0.202</td>
790
+ </tr>
791
+ <tr>
792
+ <td><strong>dots.ocr</strong> </td>
793
+ <td><strong>0.177</strong></td>
794
+ <td><strong>0.075</strong></td>
795
+ <td><strong>0.297</strong></td>
796
+ <td><strong>79.2</strong></td>
797
+ <td><strong>0.186</strong></td>
798
+ <td><strong>0.152</strong></td>
799
+ </tr>
800
+
801
+ </tbody>
802
+ </table>
803
+
804
+ > **Notes:**
805
+ > - We use the same metric calculation pipeline of [OmniDocBench](https://github.com/opendatalab/OmniDocBench).
806
+ > - We delete the Page-header and Page-footer cells in the result markdown.
807
+
808
+ #### Layout Detection
809
+
810
+ <table>
811
+ <thead>
812
+ <tr>
813
+ <th rowspan="2"><strong>Method</strong></th>
814
+ <th colspan="5" style="text-align: center;"><strong>F1@IoU=.50:.05:.95↑</strong></th>
815
+ <th colspan="5" style="text-align: center;"><strong>F1@IoU=.50↑</strong></th>
816
+ </tr>
817
+ <tr>
818
+ <th>Overall</th>
819
+ <th>Text</th>
820
+ <th>Formula</th>
821
+ <th>Table</th>
822
+ <th>Picture</th>
823
+ <th>Overall</th>
824
+ <th>Text</th>
825
+ <th>Formula</th>
826
+ <th>Table</th>
827
+ <th>Picture</th>
828
+ </tr>
829
+ </thead>
830
+
831
+ <tbody>
832
+ <td>DocLayout-YOLO-DocStructBench</td>
833
+ <td>0.733</td>
834
+ <td>0.694</td>
835
+ <td>0.480</td>
836
+ <td>0.803</td>
837
+ <td>0.619</td>
838
+ <td>0.806</td>
839
+ <td>0.779</td>
840
+ <td>0.620</td>
841
+ <td>0.858</td>
842
+ <td>0.678</td>
843
+ </tr>
844
+
845
+ <tr>
846
+ <td>dots.ocr-parse all</td>
847
+ <td>0.831</td>
848
+ <td>0.801</td>
849
+ <td>0.654</td>
850
+ <td>0.838</td>
851
+ <td>0.748</td>
852
+ <td>0.922</td>
853
+ <td>0.909</td>
854
+ <td>0.770</td>
855
+ <td>0.888</td>
856
+ <td>0.831</td>
857
+ </tr>
858
+
859
+ <tr>
860
+ <td> <strong>dots.ocr-detection only</strong> </td>
861
+ <td><strong>0.845</strong></td>
862
+ <td><strong>0.816</strong></td>
863
+ <td><strong>0.716</strong></td>
864
+ <td><strong>0.875</strong></td>
865
+ <td><strong>0.765</strong></td>
866
+ <td><strong>0.930</strong></td>
867
+ <td><strong>0.917</strong></td>
868
+ <td><strong>0.832</strong></td>
869
+ <td><strong>0.918</strong></td>
870
+ <td><strong>0.843</strong></td>
871
+ </tr>
872
+
873
+ </tbody>
874
+ </table>
875
+
876
+ > **Notes:**
877
+ > - prompt_layout_all_en for **parse all**, prompt_layout_only_en for **detection only**, please refer to [prompts](https://github.com/rednote-hilab/dots.ocr/blob/master/dots_ocr/utils/prompts.py)
878
+
879
+
880
+ ### 3. olmOCR-bench.
881
+
882
+ <table>
883
+ <thead>
884
+ <tr>
885
+ <th>Model</th>
886
+ <th>ArXiv</th>
887
+ <th>Old Scans<br>Math</th>
888
+ <th>Tables</th>
889
+ <th>Old Scans</th>
890
+ <th>Headers and<br>Footers</th>
891
+ <th>Multi<br>column</th>
892
+ <th>Long Tiny<br>Text</th>
893
+ <th>Base</th>
894
+ <th>Overall</th>
895
+ </tr>
896
+ </thead>
897
+ <tbody>
898
+ <tr>
899
+ <td>GOT OCR</td>
900
+ <td>52.7</td>
901
+ <td>52.0</td>
902
+ <td>0.2</td>
903
+ <td>22.1</td>
904
+ <td>93.6</td>
905
+ <td>42.0</td>
906
+ <td>29.9</td>
907
+ <td>94.0</td>
908
+ <td>48.3 ± 1.1</td>
909
+ </tr>
910
+ <tr>
911
+ <td>Marker</td>
912
+ <td>76.0</td>
913
+ <td>57.9</td>
914
+ <td>57.6</td>
915
+ <td>27.8</td>
916
+ <td>84.9</td>
917
+ <td>72.9</td>
918
+ <td>84.6</td>
919
+ <td>99.1</td>
920
+ <td>70.1 ± 1.1</td>
921
+ </tr>
922
+ <tr>
923
+ <td>MinerU</td>
924
+ <td>75.4</td>
925
+ <td>47.4</td>
926
+ <td>60.9</td>
927
+ <td>17.3</td>
928
+ <td><strong>96.6</strong></td>
929
+ <td>59.0</td>
930
+ <td>39.1</td>
931
+ <td>96.6</td>
932
+ <td>61.5 ± 1.1</td>
933
+ </tr>
934
+ <tr>
935
+ <td>Mistral OCR</td>
936
+ <td>77.2</td>
937
+ <td>67.5</td>
938
+ <td>60.6</td>
939
+ <td>29.3</td>
940
+ <td>93.6</td>
941
+ <td>71.3</td>
942
+ <td>77.1</td>
943
+ <td>99.4</td>
944
+ <td>72.0 ± 1.1</td>
945
+ </tr>
946
+ <tr>
947
+ <td>Nanonets OCR</td>
948
+ <td>67.0</td>
949
+ <td>68.6</td>
950
+ <td>77.7</td>
951
+ <td>39.5</td>
952
+ <td>40.7</td>
953
+ <td>69.9</td>
954
+ <td>53.4</td>
955
+ <td>99.3</td>
956
+ <td>64.5 ± 1.1</td>
957
+ </tr>
958
+ <tr>
959
+ <td>GPT-4o<br>(No Anchor)</td>
960
+ <td>51.5</td>
961
+ <td><strong>75.5</strong></td>
962
+ <td>69.1</td>
963
+ <td>40.9</td>
964
+ <td>94.2</td>
965
+ <td>68.9</td>
966
+ <td>54.1</td>
967
+ <td>96.7</td>
968
+ <td>68.9 ± 1.1</td>
969
+ </tr>
970
+ <tr>
971
+ <td>GPT-4o<br>(Anchored)</td>
972
+ <td>53.5</td>
973
+ <td>74.5</td>
974
+ <td>70.0</td>
975
+ <td>40.7</td>
976
+ <td>93.8</td>
977
+ <td>69.3</td>
978
+ <td>60.6</td>
979
+ <td>96.8</td>
980
+ <td>69.9 ± 1.1</td>
981
+ </tr>
982
+ <tr>
983
+ <td>Gemini Flash 2<br>(No Anchor)</td>
984
+ <td>32.1</td>
985
+ <td>56.3</td>
986
+ <td>61.4</td>
987
+ <td>27.8</td>
988
+ <td>48.0</td>
989
+ <td>58.7</td>
990
+ <td><strong>84.4</strong></td>
991
+ <td>94.0</td>
992
+ <td>57.8 ± 1.1</td>
993
+ </tr>
994
+ <tr>
995
+ <td>Gemini Flash 2<br>(Anchored)</td>
996
+ <td>54.5</td>
997
+ <td>56.1</td>
998
+ <td>72.1</td>
999
+ <td>34.2</td>
1000
+ <td>64.7</td>
1001
+ <td>61.5</td>
1002
+ <td>71.5</td>
1003
+ <td>95.6</td>
1004
+ <td>63.8 ± 1.2</td>
1005
+ </tr>
1006
+ <tr>
1007
+ <td>Qwen 2 VL<br>(No Anchor)</td>
1008
+ <td>19.7</td>
1009
+ <td>31.7</td>
1010
+ <td>24.2</td>
1011
+ <td>17.1</td>
1012
+ <td>88.9</td>
1013
+ <td>8.3</td>
1014
+ <td>6.8</td>
1015
+ <td>55.5</td>
1016
+ <td>31.5 ± 0.9</td>
1017
+ </tr>
1018
+ <tr>
1019
+ <td>Qwen 2.5 VL<br>(No Anchor)</td>
1020
+ <td>63.1</td>
1021
+ <td>65.7</td>
1022
+ <td>67.3</td>
1023
+ <td>38.6</td>
1024
+ <td>73.6</td>
1025
+ <td>68.3</td>
1026
+ <td>49.1</td>
1027
+ <td>98.3</td>
1028
+ <td>65.5 ± 1.2</td>
1029
+ </tr>
1030
+ <tr>
1031
+ <td>olmOCR v0.1.75<br>(No Anchor)</td>
1032
+ <td>71.5</td>
1033
+ <td>71.4</td>
1034
+ <td>71.4</td>
1035
+ <td><strong>42.8</strong></td>
1036
+ <td>94.1</td>
1037
+ <td>77.7</td>
1038
+ <td>71.0</td>
1039
+ <td>97.8</td>
1040
+ <td>74.7 ± 1.1</td>
1041
+ </tr>
1042
+ <tr>
1043
+ <td>olmOCR v0.1.75<br>(Anchored)</td>
1044
+ <td>74.9</td>
1045
+ <td>71.2</td>
1046
+ <td>71.0</td>
1047
+ <td>42.2</td>
1048
+ <td>94.5</td>
1049
+ <td>78.3</td>
1050
+ <td>73.3</td>
1051
+ <td>98.3</td>
1052
+ <td>75.5 ± 1.0</td>
1053
+ </tr>
1054
+ <tr>
1055
+ <td>MonkeyOCR-pro-3B</td>
1056
+ <td><strong>83.8</strong></td>
1057
+ <td>68.8</td>
1058
+ <td>74.6</td>
1059
+ <td>36.1</td>
1060
+ <td>91.2</td>
1061
+ <td>76.6</td>
1062
+ <td>80.1</td>
1063
+ <td>95.3</td>
1064
+ <td>75.8 ± 1.0</td>
1065
+ </tr>
1066
+ <tr>
1067
+ <td><strong>dots.ocr</strong></td>
1068
+ <td>82.1</td>
1069
+ <td>64.2</td>
1070
+ <td><strong>88.3</strong></td>
1071
+ <td>40.9</td>
1072
+ <td>94.1</td>
1073
+ <td><strong>82.4</strong></td>
1074
+ <td>81.2</td>
1075
+ <td><strong>99.5</strong></td>
1076
+ <td><strong>79.1 ± 1.0</strong></td>
1077
+ </tr>
1078
+ </tbody>
1079
+ </table>
1080
+
1081
+
1082
+ > **Note:**
1083
+ > - The metrics are from [MonkeyOCR](https://github.com/Yuliang-Liu/MonkeyOCR),
1084
+ [olmocr](https://github.com/allenai/olmocr), and our own internal evaluations.
1085
+ > - We delete the Page-header and Page-footer cells in the result markdown.
1086
+
1087
+
1088
+
1089
+ # Quick Start
1090
+ ## 1. Installation
1091
+ ### Install dots.ocr
1092
+ ```shell
1093
+ conda create -n dots_ocr python=3.12
1094
+ conda activate dots_ocr
1095
+
1096
+ git clone https://github.com/rednote-hilab/dots.ocr.git
1097
+ cd dots.ocr
1098
+
1099
+ # Install pytorch, see https://pytorch.org/get-started/previous-versions/ for your cuda version
1100
+ pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
1101
+ pip install -e .
1102
+ ```
1103
+
1104
+ If you have trouble with the installation, try our [Docker Image](https://hub.docker.com/r/rednotehilab/dots.ocr) for an easier setup, and follow these steps:
1105
+ ```shell
1106
+ git clone https://github.com/rednote-hilab/dots.ocr.git
1107
+ cd dots.ocr
1108
+ pip install -e .
1109
+ ```
1110
+
1111
+
1112
+ ### Download Model Weights
1113
+ > 💡**Note:** Please use a directory name without periods (e.g., `DotsOCR` instead of `dots.ocr`) for the model save path. This is a temporary workaround pending our integration with Transformers.
1114
+ ```shell
1115
+ python3 tools/download_model.py
1116
+ ```
1117
+
1118
+
1119
+ ## 2. Deployment
1120
+ ### vLLM inference
1121
+ We highly recommend using vllm for deployment and inference. All of our evaluations results are based on vllm version 0.9.1.
1122
+ The [Docker Image](https://hub.docker.com/r/rednotehilab/dots.ocr) is based on the official vllm image. You can also follow [Dockerfile](https://github.com/rednote-hilab/dots.ocr/blob/master/docker/Dockerfile) to build the deployment environment by yourself.
1123
+
1124
+ ```shell
1125
+ # You need to register model to vllm at first
1126
+ python3 tools/download_model.py
1127
+ export hf_model_path=./weights/DotsOCR # Path to your downloaded model weights, Please use a directory name without periods (e.g., `DotsOCR` instead of `dots.ocr`) for the model save path. This is a temporary workaround pending our integration with Transformers.
1128
+ export PYTHONPATH=$(dirname "$hf_model_path"):$PYTHONPATH
1129
+ sed -i '/^from vllm\.entrypoints\.cli\.main import main$/a\
1130
+ from DotsOCR import modeling_dots_ocr_vllm' `which vllm` # If you downloaded model weights by yourself, please replace `DotsOCR` by your model saved directory name, and remember to use a directory name without periods (e.g., `DotsOCR` instead of `dots.ocr`)
1131
+
1132
+ # launch vllm server
1133
+ CUDA_VISIBLE_DEVICES=0 vllm serve ${hf_model_path} --tensor-parallel-size 1 --gpu-memory-utilization 0.95 --chat-template-content-format string --served-model-name model --trust-remote-code
1134
+
1135
+ # If you get a ModuleNotFoundError: No module named 'DotsOCR', please check the note above on the saved model directory name.
1136
+
1137
+ # vllm api demo
1138
+ python3 ./demo/demo_vllm.py --prompt_mode prompt_layout_all_en
1139
+ ```
1140
+
1141
+ ### Hugginface inference
1142
+ ```shell
1143
+ python3 demo/demo_hf.py
1144
+ ```
1145
+
1146
+ <details>
1147
+ <summary><b>Hugginface inference details</b></summary>
1148
+
1149
+ ```python
1150
+ import torch
1151
+ from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
1152
+ from qwen_vl_utils import process_vision_info
1153
+ from dots_ocr.utils import dict_promptmode_to_prompt
1154
+
1155
+ model_path = "./weights/DotsOCR"
1156
+ model = AutoModelForCausalLM.from_pretrained(
1157
+ model_path,
1158
+ attn_implementation="flash_attention_2",
1159
+ torch_dtype=torch.bfloat16,
1160
+ device_map="auto",
1161
+ trust_remote_code=True
1162
+ )
1163
+ processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
1164
+
1165
+ image_path = "demo/demo_image1.jpg"
1166
+ prompt = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
1167
+
1168
+ 1. Bbox format: [x1, y1, x2, y2]
1169
+
1170
+ 2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
1171
+
1172
+ 3. Text Extraction & Formatting Rules:
1173
+ - Picture: For the 'Picture' category, the text field should be omitted.
1174
+ - Formula: Format its text as LaTeX.
1175
+ - Table: Format its text as HTML.
1176
+ - All Others (Text, Title, etc.): Format their text as Markdown.
1177
+
1178
+ 4. Constraints:
1179
+ - The output text must be the original text from the image, with no translation.
1180
+ - All layout elements must be sorted according to human reading order.
1181
+
1182
+ 5. Final Output: The entire output must be a single JSON object.
1183
+ """
1184
+
1185
+ messages = [
1186
+ {
1187
+ "role": "user",
1188
+ "content": [
1189
+ {
1190
+ "type": "image",
1191
+ "image": image_path
1192
+ },
1193
+ {"type": "text", "text": prompt}
1194
+ ]
1195
+ }
1196
+ ]
1197
+
1198
+ # Preparation for inference
1199
+ text = processor.apply_chat_template(
1200
+ messages,
1201
+ tokenize=False,
1202
+ add_generation_prompt=True
1203
+ )
1204
+ image_inputs, video_inputs = process_vision_info(messages)
1205
+ inputs = processor(
1206
+ text=[text],
1207
+ images=image_inputs,
1208
+ videos=video_inputs,
1209
+ padding=True,
1210
+ return_tensors="pt",
1211
+ )
1212
+
1213
+ inputs = inputs.to("cuda")
1214
+
1215
+ # Inference: Generation of the output
1216
+ generated_ids = model.generate(**inputs, max_new_tokens=24000)
1217
+ generated_ids_trimmed = [
1218
+ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
1219
+ ]
1220
+ output_text = processor.batch_decode(
1221
+ generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
1222
+ )
1223
+ print(output_text)
1224
+
1225
+ ```
1226
+
1227
+ </details>
1228
+
1229
+ ## 3. Document Parse
1230
+ **Based on vLLM server**, you can parse an image or a pdf file using the following commands:
1231
+ ```bash
1232
+
1233
+ # Parse all layout info, both detection and recognition
1234
+ # Parse a single image
1235
+ python3 dots_ocr/parser.py demo/demo_image1.jpg
1236
+ # Parse a single PDF
1237
+ python3 dots_ocr/parser.py demo/demo_pdf1.pdf --num_threads 64 # try bigger num_threads for pdf with a large number of pages
1238
+
1239
+ # Layout detection only
1240
+ python3 dots_ocr/parser.py demo/demo_image1.jpg --prompt prompt_layout_only_en
1241
+
1242
+ # Parse text only, except Page-header and Page-footer
1243
+ python3 dots_ocr/parser.py demo/demo_image1.jpg --prompt prompt_ocr
1244
+
1245
+ # Parse layout info by bbox
1246
+ python3 dots_ocr/parser.py demo/demo_image1.jpg --prompt prompt_grounding_ocr --bbox 163 241 1536 705
1247
+
1248
+ ```
1249
+
1250
+ <details>
1251
+ <summary><b>Output Results</b></summary>
1252
+
1253
+ 1. **Structured Layout Data** (`demo_image1.json`): A JSON file containing the detected layout elements, including their bounding boxes, categories, and extracted text.
1254
+ 2. **Processed Markdown File** (`demo_image1.md`): A Markdown file generated from the concatenated text of all detected cells.
1255
+ * An additional version, `demo_image1_nohf.md`, is also provided, which excludes page headers and footers for compatibility with benchmarks like Omnidocbench and olmOCR-bench.
1256
+ 3. **Layout Visualization** (`demo_image1.jpg`): The original image with the detected layout bounding boxes drawn on it.
1257
+
1258
+ </details>
1259
+
1260
+ ## 4. Demo
1261
+ You can run the demo with the following command, or try directly at [live demo](https://dotsocr.xiaohongshu.com/)
1262
+ ```bash
1263
+ python demo/demo_gradio.py
1264
+ ```
1265
+
1266
+ We also provide a demo for grounding ocr:
1267
+ ```bash
1268
+ python demo/demo_gradio_annotion.py
1269
+ ```
1270
+
1271
+
1272
+ ### Example for formula document
1273
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/formula1.png" alt="formula1.png" border="0" />
1274
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/formula2.png" alt="formula2.png" border="0" />
1275
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/formula3.png" alt="formula3.png" border="0" />
1276
+
1277
+ ### Example for table document
1278
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/table1.png" alt="table1.png" border="0" />
1279
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/table2.png" alt="table2.png" border="0" />
1280
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/table3.png" alt="table3.png" border="0" />
1281
+
1282
+ ### Example for multilingual document
1283
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/Tibetan.png" alt="Tibetan.png" border="0" />
1284
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/tradition_zh.png" alt="tradition_zh.png" border="0" />
1285
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/nl.png" alt="nl.png" border="0" />
1286
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/kannada.png" alt="kannada.png" border="0" />
1287
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/russian.png" alt="russian.png" border="0" />
1288
+
1289
+ ### Example for reading order
1290
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/reading_order.png" alt="reading_order.png" border="0" />
1291
+
1292
+ ### Example for grounding ocr
1293
+ <img src="https://raw.githubusercontent.com/rednote-hilab/dots.ocr/master/assets/showcase/grounding.png" alt="grounding.png" border="0" />
1294
+
1295
+
1296
+ ## Acknowledgments
1297
+ We would like to thank [Qwen2.5-VL](https://github.com/QwenLM/Qwen2.5-VL), [aimv2](https://github.com/apple/ml-aim), [MonkeyOCR](https://github.com/Yuliang-Liu/MonkeyOCR),
1298
+ [OmniDocBench](https://github.com/opendatalab/OmniDocBench), [PyMuPDF](https://github.com/pymupdf/PyMuPDF), for providing code and models.
1299
+
1300
+ We also thank [DocLayNet](https://github.com/DS4SD/DocLayNet), [M6Doc](https://github.com/HCIILAB/M6Doc), [CDLA](https://github.com/buptlihang/CDLA), [D4LA](https://github.com/AlibabaResearch/AdvancedLiterateMachinery) for providing valuable datasets.
1301
+
1302
+ ## Limitation & Future Work
1303
+
1304
+ - **Complex Document Elements:**
1305
+ - **Table&Formula**: dots.ocr is not yet perfect for high-complexity tables and formula extraction.
1306
+ - **Picture**: Pictures in documents are currently not parsed.
1307
+
1308
+ - **Parsing Failures:** The model may fail to parse under certain conditions:
1309
+ - When the character-to-pixel ratio is excessively high. Try enlarging the image or increasing the PDF parsing DPI (a setting of 200 is recommended). However, please note that the model performs optimally on images with a resolution under 11289600 pixels.
1310
+ - Continuous special characters, such as ellipses (`...`) and underscores (`_`), may cause the prediction output to repeat endlessly. In such scenarios, consider using alternative prompts like `prompt_layout_only_en`, `prompt_ocr`, or `prompt_grounding_ocr` ([details here](https://github.com/rednote-hilab/dots.ocr/blob/master/dots_ocr/utils/prompts.py)).
1311
+
1312
+ - **Performance Bottleneck:** Despite its 1.7B parameter LLM foundation, **dots.ocr** is not yet optimized for high-throughput processing of large PDF volumes.
1313
+
1314
+ We are committed to achieving more accurate table and formula parsing, as well as enhancing the model's OCR capabilities for broader generalization, all while aiming for **a more powerful, more efficient model**. Furthermore, we are actively considering the development of **a more general-purpose perception model** based on Vision-Language Models (VLMs), which would integrate general detection, image captioning, and OCR tasks into a unified framework. **Parsing the content of the pictures in the documents** is also a key priority for our future work.
1315
+ We believe that collaboration is the key to tackling these exciting challenges. If you are passionate about advancing the frontiers of document intelligence and are interested in contributing to these future endeavors, we would love to hear from you. Please reach out to us via email at: [yanqing4@xiaohongshu.com].
chat_template.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{%- for m in messages %}{%- if m.role == 'system' %}{{- '<|system|>' + m.content + '<|endofsystem|>\n' }}{%- elif m.role == 'user' %}{% if m.content is string %}{{- '<|user|>' + m.content + '<|endofuser|>' }}{% else %} {% for content in m.content %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|img|><|imgpad|><|endofimg|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|img|><|video_pad|><|endofimg|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{%- endif %}{%- elif m.role == 'assistant' %}{{- '<|assistant|>' + m.content }}{%- if not loop.last %}{{- '<|endofassistant|>' }}{%- endif %}{%- endif %}{%- endfor %}{%- if messages[-1].role != 'assistant' %}{{- '<|assistant|>' }}{%- endif %}"
3
+ }
config.json ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "DotsOCRForCausalLM"
4
+ ],
5
+ "model_type": "dots_ocr",
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_dots.DotsOCRConfig",
8
+ "AutoModelForCausalLM": "modeling_dots_ocr.DotsOCRForCausalLM"
9
+ },
10
+ "attention_bias": true,
11
+ "attention_dropout": 0.0,
12
+ "hidden_act": "silu",
13
+ "hidden_size": 1536,
14
+ "initializer_range": 0.02,
15
+ "intermediate_size": 8960,
16
+ "max_position_embeddings": 131072,
17
+ "max_window_layers": 28,
18
+ "num_attention_heads": 12,
19
+ "num_hidden_layers": 28,
20
+ "num_key_value_heads": 2,
21
+ "rms_norm_eps": 1e-06,
22
+ "rope_scaling": null,
23
+ "rope_theta": 1000000,
24
+ "sliding_window": 131072,
25
+ "tie_word_embeddings": false,
26
+ "torch_dtype": "bfloat16",
27
+ "transformers_version": "4.51.0",
28
+ "use_cache": true,
29
+ "use_sliding_window": false,
30
+ "vocab_size": 151936,
31
+ "image_token_id": 151665,
32
+ "video_token_id": 151656,
33
+ "vision_config": {
34
+ "embed_dim": 1536,
35
+ "hidden_size": 1536,
36
+ "intermediate_size": 4224,
37
+ "num_hidden_layers": 42,
38
+ "num_attention_heads": 12,
39
+ "num_channels": 3,
40
+ "patch_size": 14,
41
+ "post_norm": true,
42
+ "rms_norm_eps": 1e-05,
43
+ "spatial_merge_size": 2,
44
+ "temporal_patch_size": 1,
45
+ "use_bias": false,
46
+ "attn_implementation": "flash_attention_2",
47
+ "init_merger_std": 0.02,
48
+ "initializer_range": 0.02,
49
+ "is_causal": false
50
+ },
51
+ "eos_token_id": [
52
+ 151643,
53
+ 151673
54
+ ],
55
+ "quantization": {
56
+ "group_size": 64,
57
+ "bits": 8,
58
+ "mode": "affine"
59
+ },
60
+ "quantization_config": {
61
+ "group_size": 64,
62
+ "bits": 8,
63
+ "mode": "affine"
64
+ }
65
+ }
configuration_dots.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Optional
2
+ from transformers.configuration_utils import PretrainedConfig
3
+ from transformers.models.qwen2 import Qwen2Config
4
+ from transformers import Qwen2_5_VLProcessor, AutoProcessor
5
+ from transformers.models.auto.configuration_auto import CONFIG_MAPPING
6
+
7
+
8
+ class DotsVisionConfig(PretrainedConfig):
9
+ model_type: str = "dots_vit"
10
+
11
+ def __init__(
12
+ self,
13
+ embed_dim: int = 1536, # vision encoder embed size
14
+ hidden_size: int = 1536, # after merger hidden size
15
+ intermediate_size: int = 4224,
16
+ num_hidden_layers: int = 42,
17
+ num_attention_heads: int = 12,
18
+ num_channels: int = 3,
19
+ patch_size: int = 14,
20
+ spatial_merge_size: int = 2,
21
+ temporal_patch_size: int = 1,
22
+ rms_norm_eps: float = 1e-5,
23
+ use_bias: bool = False,
24
+ attn_implementation="flash_attention_2", # "eager","sdpa","flash_attention_2"
25
+ initializer_range=0.02,
26
+ init_merger_std=0.02,
27
+ is_causal=False, # ve causal forward
28
+ post_norm=True,
29
+ gradient_checkpointing=False,
30
+ **kwargs: Any,
31
+ ):
32
+ super().__init__(**kwargs)
33
+ self.embed_dim = embed_dim
34
+ self.hidden_size = hidden_size
35
+ self.intermediate_size = intermediate_size
36
+ self.num_hidden_layers = num_hidden_layers
37
+ self.num_attention_heads = num_attention_heads
38
+ self.num_channels = num_channels
39
+ self.patch_size = patch_size
40
+ self.spatial_merge_size = spatial_merge_size
41
+ self.temporal_patch_size = temporal_patch_size
42
+ self.rms_norm_eps = rms_norm_eps
43
+ self.use_bias = use_bias
44
+ self.attn_implementation = attn_implementation
45
+ self.initializer_range = initializer_range
46
+ self.init_merger_std = init_merger_std
47
+ self.is_causal = is_causal
48
+ self.post_norm = post_norm
49
+ self.gradient_checkpointing = gradient_checkpointing
50
+
51
+
52
+
53
+ class DotsOCRConfig(Qwen2Config):
54
+ model_type = "dots_ocr"
55
+ def __init__(self,
56
+ image_token_id = 151665,
57
+ video_token_id = 151656,
58
+ vision_config: Optional[dict] = None, *args, **kwargs):
59
+ super().__init__(*args, **kwargs)
60
+ self.image_token_id = image_token_id
61
+ self.video_token_id = video_token_id
62
+ self.vision_config = DotsVisionConfig(**(vision_config or {}))
63
+
64
+ def save_pretrained(self, save_directory, **kwargs):
65
+ self._auto_class = None
66
+ super().save_pretrained(save_directory, **kwargs)
67
+
68
+
69
+ class DotsVLProcessor(Qwen2_5_VLProcessor):
70
+ def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
71
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
72
+ self.image_token = "<|imgpad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
73
+ self.image_token_id = 151665 if not hasattr(tokenizer, "image_token_id") else tokenizer.image_token_id
74
+
75
+
76
+ AutoProcessor.register("dots_ocr", DotsVLProcessor)
77
+ CONFIG_MAPPING.register("dots_ocr", DotsOCRConfig)
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "max_length": 32768,
3
+ "eos_token_id": [
4
+ 151643,
5
+ 151673
6
+ ]
7
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:989b3801777747df653d23ac0bf14c2a512dd8720e917dda8fb3fb510e25e873
3
+ size 4412592228
modeling_dots_ocr.py ADDED
@@ -0,0 +1,131 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Optional, Tuple, Union
2
+
3
+ import torch
4
+ from transformers.modeling_outputs import CausalLMOutputWithPast
5
+ from transformers.models.qwen2 import Qwen2ForCausalLM
6
+
7
+ from .configuration_dots import DotsVisionConfig, DotsOCRConfig
8
+ from .modeling_dots_vision import DotsVisionTransformer
9
+
10
+
11
+ DOTS_VLM_MAX_IMAGES = 200
12
+
13
+
14
+ class DotsOCRForCausalLM(Qwen2ForCausalLM):
15
+ config_class = DotsOCRConfig
16
+
17
+ def __init__(self, config: DotsOCRConfig):
18
+ super().__init__(config)
19
+
20
+ if isinstance(self.config.vision_config, dict):
21
+ vision_config = DotsVisionConfig(**self.config.vision_config)
22
+ self.config.vision_config = vision_config
23
+ else:
24
+ vision_config = self.config.vision_config
25
+
26
+ self.vision_tower = DotsVisionTransformer(vision_config)
27
+
28
+ def prepare_inputs_embeds(
29
+ self,
30
+ input_ids: torch.LongTensor,
31
+ pixel_values: Optional[torch.FloatTensor] = None,
32
+ grid_thw: Optional[torch.FloatTensor] = None,
33
+ img_mask: Optional[torch.BoolTensor] = None,
34
+ ) -> torch.Tensor:
35
+ inputs_embeds = self.get_input_embeddings()(input_ids)
36
+
37
+ if pixel_values is not None:
38
+ assert img_mask is not None
39
+ if grid_thw.shape[0] > DOTS_VLM_MAX_IMAGES:
40
+ print(
41
+ f"Num image exceeded: {grid_thw.shape[0]} > {DOTS_VLM_MAX_IMAGES}, which may cause FSDP hang"
42
+ )
43
+
44
+ vision_embeddings = self.vision_tower(pixel_values, grid_thw)
45
+
46
+ true_indices = torch.nonzero(img_mask).squeeze()
47
+ if len(true_indices) > vision_embeddings.size(0):
48
+ print(
49
+ f"img_mask sum > VE and will be truncated, mask.sum()={len(true_indices)} {vision_embeddings.size(0)=}"
50
+ )
51
+ true_indices = true_indices[: vision_embeddings.size(0)]
52
+ new_img_mask = torch.zeros_like(img_mask, device=img_mask.device)
53
+ new_img_mask[true_indices[:, 0], true_indices[:, 1]] = True
54
+ else:
55
+ new_img_mask = img_mask
56
+
57
+ assert (
58
+ vision_embeddings.size(0) == new_img_mask.sum()
59
+ ), f"{vision_embeddings.size(0)=}, {new_img_mask.sum()=}"
60
+
61
+ inputs_embeds = inputs_embeds.masked_scatter(
62
+ new_img_mask.to(inputs_embeds.device).unsqueeze(-1).expand_as(inputs_embeds),
63
+ vision_embeddings.to(inputs_embeds.device).type(inputs_embeds.dtype),
64
+ )
65
+
66
+ return inputs_embeds
67
+
68
+ def forward(
69
+ self,
70
+ input_ids: torch.LongTensor,
71
+ pixel_values: Optional[torch.FloatTensor] = None,
72
+ image_grid_thw: Optional[torch.FloatTensor] = None,
73
+ inputs_embeds: Optional[torch.Tensor] = None,
74
+ attention_mask: Optional[torch.Tensor] = None,
75
+ position_ids: Optional[torch.LongTensor] = None,
76
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
77
+ labels: Optional[torch.LongTensor] = None,
78
+ output_attentions: Optional[bool] = None,
79
+ output_hidden_states: Optional[bool] = None,
80
+ return_dict: Optional[bool] = None,
81
+ use_cache: Optional[bool] = None,
82
+ logits_to_keep: int = 0,
83
+ **loss_kwargs,
84
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
85
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
86
+ assert len(input_ids) >= 1, f"empty input_ids {input_ids.shape=} will cause gradnorm nan"
87
+ if inputs_embeds is None:
88
+ img_mask = input_ids == self.config.image_token_id
89
+ inputs_embeds = self.prepare_inputs_embeds(input_ids, pixel_values, image_grid_thw, img_mask)
90
+
91
+ outputs = super().forward(
92
+ inputs_embeds=inputs_embeds,
93
+ attention_mask=attention_mask,
94
+ position_ids=position_ids,
95
+ past_key_values=past_key_values,
96
+ labels=labels,
97
+ use_cache=use_cache if use_cache is not None else self.config.use_cache,
98
+ output_attentions=output_attentions,
99
+ output_hidden_states=output_hidden_states,
100
+ # return_dict=return_dict,
101
+ logits_to_keep=logits_to_keep,
102
+ **loss_kwargs,
103
+ )
104
+
105
+ return outputs
106
+
107
+ def prepare_inputs_for_generation(
108
+ self,
109
+ input_ids,
110
+ past_key_values=None,
111
+ inputs_embeds=None,
112
+ pixel_values=None,
113
+ attention_mask=None,
114
+ cache_position=None,
115
+ num_logits_to_keep=None,
116
+ **kwargs,
117
+ ):
118
+ model_inputs = super().prepare_inputs_for_generation(
119
+ input_ids,
120
+ past_key_values=past_key_values,
121
+ inputs_embeds=inputs_embeds,
122
+ attention_mask=attention_mask,
123
+ cache_position=cache_position,
124
+ num_logits_to_keep=num_logits_to_keep,
125
+ **kwargs,
126
+ )
127
+
128
+ if cache_position[0] == 0:
129
+ model_inputs["pixel_values"] = pixel_values
130
+
131
+ return model_inputs
modeling_dots_ocr_vllm.py ADDED
@@ -0,0 +1,451 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from functools import cached_property
2
+ from typing import Iterable, Literal, Mapping, Optional, Set, Tuple, TypedDict, Union
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ from transformers.models.qwen2_vl import Qwen2VLImageProcessor, Qwen2VLProcessor
7
+ from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
8
+ from vllm import ModelRegistry
9
+ from vllm.config import VllmConfig
10
+ from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
11
+ from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal
12
+ from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM
13
+ from vllm.model_executor.models.qwen2_5_vl import (
14
+ Qwen2_5_VLMultiModalProcessor,
15
+ Qwen2_5_VLProcessingInfo,
16
+ )
17
+ from vllm.model_executor.models.qwen2_vl import Qwen2VLDummyInputsBuilder
18
+ from vllm.model_executor.models.utils import (
19
+ AutoWeightsLoader,
20
+ WeightsMapper,
21
+ init_vllm_registered_model,
22
+ maybe_prefix,
23
+ merge_multimodal_embeddings,
24
+ )
25
+ from vllm.model_executor.sampling_metadata import SamplingMetadata
26
+ from vllm.multimodal import MULTIMODAL_REGISTRY
27
+ from vllm.multimodal.inputs import MultiModalDataDict
28
+ from vllm.multimodal.parse import ImageSize
29
+ from vllm.sequence import IntermediateTensors
30
+
31
+ from .configuration_dots import DotsVisionConfig
32
+ from .configuration_dots import DotsOCRConfig
33
+ from .modeling_dots_vision import DotsVisionTransformer
34
+
35
+
36
+ class DotsOCRImagePixelInputs(TypedDict):
37
+ type: Literal["pixel_values", "image_grid_thw"]
38
+
39
+ pixel_values: torch.Tensor
40
+ image_grid_thw: torch.Tensor
41
+
42
+
43
+ class DotsOCRImageEmbeddingInputs(TypedDict):
44
+ type: Literal["image_embeds", "image_grid_thw"]
45
+ image_embeds: torch.Tensor
46
+ """Supported types:
47
+ - List[`torch.Tensor`]: A list of tensors holding all images' features.
48
+ Each tensor holds an image's features.
49
+ - `torch.Tensor`: A tensor holding all images' features
50
+ (concatenation of all images' feature tensors).
51
+
52
+ Tensor shape: `(num_image_features, hidden_size)`
53
+ - `num_image_features` varies based on
54
+ the number and resolution of the images.
55
+ - `hidden_size` must match the hidden size of language model backbone.
56
+ """
57
+
58
+ image_grid_thw: torch.Tensor
59
+
60
+
61
+ DotsOCRImageInputs = Union[DotsOCRImagePixelInputs, DotsOCRImageEmbeddingInputs]
62
+
63
+
64
+ class DotsOCRMultiModalProcessor(Qwen2_5_VLMultiModalProcessor):
65
+ pass
66
+
67
+
68
+ class DotsOCRDummyInputsBuilder(Qwen2VLDummyInputsBuilder):
69
+ def get_dummy_mm_data(
70
+ self,
71
+ seq_len: int,
72
+ mm_counts: Mapping[str, int],
73
+ ) -> MultiModalDataDict:
74
+ num_images = mm_counts.get("image", 0)
75
+
76
+ target_width, target_height = self.info.get_image_size_with_most_features()
77
+
78
+ return {
79
+ "image": self._get_dummy_images(width=target_width, height=target_height, num_images=num_images),
80
+ }
81
+
82
+
83
+ class DotsOCRProcessingInfo(Qwen2_5_VLProcessingInfo):
84
+ def get_hf_config(self) -> DotsOCRConfig:
85
+ config = self.ctx.get_hf_config()
86
+ if not config.__class__.__name__ == 'DotsOCRConfig':
87
+ raise TypeError(f"Expected DotsOCRConfig, got {type(config)}")
88
+
89
+ if hasattr(config, "vision_config") and isinstance(config.vision_config, dict):
90
+ config.vision_config = DotsVisionConfig(**config.vision_config)
91
+
92
+ return config
93
+
94
+ def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
95
+ return {"image": None, "video": 0}
96
+
97
+ def get_mm_max_tokens_per_item(
98
+ self,
99
+ seq_len: int,
100
+ mm_counts: Mapping[str, int],
101
+ ) -> Mapping[str, int]:
102
+ max_image_tokens = self.get_max_image_tokens()
103
+ return {"image": max_image_tokens, "video": 0}
104
+
105
+ def get_hf_processor(
106
+ self,
107
+ *,
108
+ min_pixels: Optional[int] = None,
109
+ max_pixels: Optional[int] = None,
110
+ size: Optional[dict[str, int]] = None,
111
+ **kwargs: object,
112
+ ) -> Qwen2VLProcessor:
113
+ self.get_tokenizer().image_token = "<|imgpad|>" # Ensure image token is set
114
+ processor = self.ctx.get_hf_processor(
115
+ Qwen2VLProcessor,
116
+ image_processor=self.get_image_processor(min_pixels=min_pixels, max_pixels=max_pixels, size=size),
117
+ **kwargs,
118
+ )
119
+ processor.image_token = "<|imgpad|>"
120
+ processor.video_token = "<|video_pad|>"
121
+ return processor
122
+
123
+ def _get_vision_info(
124
+ self,
125
+ *,
126
+ image_width: int,
127
+ image_height: int,
128
+ num_frames: int = 1,
129
+ do_resize: bool = True,
130
+ image_processor: Optional[Qwen2VLImageProcessor],
131
+ ) -> tuple[ImageSize, int]:
132
+ if image_processor is None:
133
+ image_processor = self.get_image_processor()
134
+
135
+ hf_config: DotsOCRConfig = self.get_hf_config()
136
+ vision_config = hf_config.vision_config
137
+ patch_size = vision_config.patch_size
138
+ merge_size = vision_config.spatial_merge_size
139
+ temporal_patch_size = vision_config.temporal_patch_size
140
+
141
+ if do_resize:
142
+ resized_height, resized_width = smart_resize(
143
+ height=image_height,
144
+ width=image_width,
145
+ factor=patch_size * merge_size,
146
+ min_pixels=image_processor.min_pixels,
147
+ max_pixels=image_processor.max_pixels,
148
+ )
149
+ preprocessed_size = ImageSize(width=resized_width, height=resized_height)
150
+ else:
151
+ preprocessed_size = ImageSize(width=image_width, height=image_height)
152
+
153
+ # NOTE: Frames are padded to be divisible by `temporal_patch_size`
154
+ # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294
155
+ padded_num_frames = num_frames + num_frames % temporal_patch_size
156
+
157
+ grid_t = max(padded_num_frames // temporal_patch_size, 1)
158
+ grid_h = preprocessed_size.height // patch_size
159
+ grid_w = preprocessed_size.width // patch_size
160
+
161
+ num_patches = grid_t * grid_h * grid_w
162
+ num_vision_tokens = num_patches // (merge_size**2)
163
+
164
+ return preprocessed_size, num_vision_tokens
165
+
166
+
167
+ @MULTIMODAL_REGISTRY.register_processor(
168
+ Qwen2_5_VLMultiModalProcessor,
169
+ info=DotsOCRProcessingInfo,
170
+ dummy_inputs=DotsOCRDummyInputsBuilder,
171
+ )
172
+ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal):
173
+ hf_to_vllm_mapper = WeightsMapper(
174
+ orig_to_new_prefix={
175
+ "lm_head.": "language_model.lm_head.",
176
+ "model.": "language_model.model.",
177
+ }
178
+ )
179
+ _tp_plan = {}
180
+
181
+ @classmethod
182
+ def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
183
+ if modality in ("image",):
184
+ return "<|img|><|imgpad|><|endofimg|>"
185
+
186
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
187
+ super().__init__()
188
+
189
+ self.config: DotsOCRConfig = vllm_config.model_config.hf_config
190
+ self.quant_config = vllm_config.quant_config
191
+ self.multimodal_config = vllm_config.model_config.multimodal_config
192
+
193
+ if isinstance(self.config.vision_config, dict):
194
+ vision_config = DotsVisionConfig(**self.config.vision_config)
195
+ self.config.vision_config = vision_config
196
+ else:
197
+ vision_config = self.config.vision_config
198
+
199
+ self.vision_tower = DotsVisionTransformer(vision_config)
200
+ self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(
201
+ vllm_config=vllm_config,
202
+ hf_config=self.config,
203
+ prefix=maybe_prefix(prefix, "language_model"),
204
+ architectures=["Qwen2ForCausalLM"],
205
+ )
206
+
207
+ @cached_property
208
+ def sampler(self):
209
+ if hasattr(self.language_model, "sampler"):
210
+ return self.language_model.sampler
211
+
212
+ return get_sampler()
213
+
214
+ def _validate_and_reshape_mm_tensor(self, mm_input: object, name: str) -> torch.Tensor:
215
+ if not isinstance(mm_input, (torch.Tensor, list)):
216
+ raise ValueError(f"Incorrect type of {name}. " f"Got type: {type(mm_input)}")
217
+ if isinstance(mm_input, torch.Tensor):
218
+ if mm_input.ndim == 2:
219
+ return mm_input
220
+ if mm_input.ndim != 3:
221
+ raise ValueError(
222
+ f"{name} should be 2D or batched 3D tensor. "
223
+ f"Got ndim: {mm_input.ndim} "
224
+ f"(shape={mm_input.shape})"
225
+ )
226
+ return torch.concat(list(mm_input))
227
+ else:
228
+ return torch.concat(mm_input)
229
+
230
+ def _parse_and_validate_image_input(self, **kwargs: object) -> Optional[DotsOCRImageInputs]:
231
+ pixel_values = kwargs.pop("pixel_values", None)
232
+ image_embeds = kwargs.pop("image_embeds", None)
233
+ image_grid_thw = kwargs.pop("image_grid_thw", None)
234
+
235
+ if pixel_values is None and image_embeds is None:
236
+ return None
237
+
238
+ if pixel_values is not None:
239
+ pixel_values = self._validate_and_reshape_mm_tensor(pixel_values, "image pixel values")
240
+ image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
241
+
242
+ if not isinstance(pixel_values, (torch.Tensor, list)):
243
+ raise ValueError("Incorrect type of image pixel values. " f"Got type: {type(pixel_values)}")
244
+
245
+ return DotsOCRImagePixelInputs(
246
+ type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw
247
+ )
248
+
249
+ if image_embeds is not None:
250
+ image_embeds = self._validate_and_reshape_mm_tensor(image_embeds, "image embeds")
251
+ image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw")
252
+
253
+ if not isinstance(image_embeds, torch.Tensor):
254
+ raise ValueError("Incorrect type of image embeddings. " f"Got type: {type(image_embeds)}")
255
+ return DotsOCRImageEmbeddingInputs(
256
+ type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw
257
+ )
258
+
259
+ def vision_forward(self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor):
260
+ from vllm.distributed import (
261
+ get_tensor_model_parallel_group,
262
+ get_tensor_model_parallel_rank,
263
+ get_tensor_model_parallel_world_size,
264
+ )
265
+
266
+ assert self.vision_tower is not None
267
+
268
+ tp_rank = get_tensor_model_parallel_rank()
269
+ tp = get_tensor_model_parallel_world_size()
270
+
271
+ image_grid_thw_chunk = image_grid_thw.chunk(tp)
272
+ image_sizes_consum = torch.tensor([i.prod(-1).sum() for i in image_grid_thw_chunk]).cumsum(dim=0)
273
+ merge_size_square = self.vision_tower.config.spatial_merge_size**2
274
+ image_embedding = torch.zeros(
275
+ (
276
+ pixel_values.shape[0] // merge_size_square,
277
+ self.vision_tower.config.hidden_size,
278
+ ),
279
+ device=pixel_values.device,
280
+ dtype=pixel_values.dtype,
281
+ )
282
+
283
+ if tp_rank < len(image_sizes_consum):
284
+ idx_start = 0 if tp_rank == 0 else image_sizes_consum[tp_rank - 1].item()
285
+ idx_end = image_sizes_consum[tp_rank].item()
286
+ pixel_values_part = pixel_values[idx_start:idx_end]
287
+ image_grid_thw_part = image_grid_thw_chunk[tp_rank]
288
+ image_embedding_part = self.vision_tower(pixel_values_part, image_grid_thw_part)
289
+ image_embedding[idx_start // merge_size_square : idx_end // merge_size_square] = image_embedding_part
290
+
291
+ group = get_tensor_model_parallel_group().device_group
292
+ torch.distributed.all_reduce(image_embedding, group=group)
293
+ return image_embedding
294
+
295
+ def _process_image_input(self, image_input: DotsOCRImageInputs) -> tuple[torch.Tensor, ...]:
296
+ grid_thw = image_input["image_grid_thw"]
297
+ assert grid_thw.ndim == 2
298
+
299
+ if image_input["type"] == "image_embeds":
300
+ image_embeds = image_input["image_embeds"].type(self.vision_tower.dtype)
301
+ else:
302
+ pixel_values = image_input["pixel_values"].type(self.vision_tower.dtype)
303
+ image_embeds = self.vision_forward(pixel_values, grid_thw)[
304
+ :, : self.config.hidden_size
305
+ ]
306
+
307
+ # Split concatenated embeddings for each image item.
308
+ merge_size = self.vision_tower.config.spatial_merge_size
309
+ sizes = grid_thw.prod(-1) // merge_size // merge_size
310
+
311
+ return image_embeds.split(sizes.tolist())
312
+
313
+ def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
314
+ modalities = {}
315
+
316
+ # Preserve the order of modalities if there are multiple of them
317
+ # from the order of kwargs.
318
+ for input_key in kwargs:
319
+ if input_key in ("pixel_values", "image_embeds") and "images" not in modalities:
320
+ modalities["images"] = self._parse_and_validate_image_input(**kwargs)
321
+ return modalities
322
+
323
+ def get_language_model(self) -> torch.nn.Module:
324
+ return self.language_model
325
+
326
+ def get_multimodal_embeddings(self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
327
+ modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
328
+ if not modalities:
329
+ return None
330
+
331
+ # The result multimodal_embeddings is tuple of tensors, with each
332
+ # tensor correspoending to a multimodal data item (image or video).
333
+ multimodal_embeddings: tuple[torch.Tensor, ...] = ()
334
+
335
+ # NOTE: It is important to iterate over the keys in this dictionary
336
+ # to preserve the order of the modalities.
337
+ for modality in modalities:
338
+ if modality == "images":
339
+ image_input = modalities["images"]
340
+ vision_embeddings = self._process_image_input(image_input)
341
+ multimodal_embeddings += vision_embeddings
342
+
343
+ return multimodal_embeddings
344
+
345
+ def get_input_embeddings(
346
+ self,
347
+ input_ids: torch.Tensor,
348
+ multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
349
+ ) -> torch.Tensor:
350
+ inputs_embeds = self.language_model.get_input_embeddings(input_ids)
351
+ if multimodal_embeddings is not None:
352
+ inputs_embeds = merge_multimodal_embeddings(
353
+ input_ids,
354
+ inputs_embeds,
355
+ multimodal_embeddings,
356
+ [self.config.image_token_id, self.config.video_token_id],
357
+ )
358
+
359
+ return inputs_embeds
360
+
361
+ def get_input_embeddings_v0(
362
+ self,
363
+ input_ids: torch.Tensor,
364
+ image_input: Optional[DotsOCRImagePixelInputs] = None,
365
+ ) -> torch.Tensor:
366
+ inputs_embeds = self.get_input_embeddings(input_ids)
367
+ if image_input is not None:
368
+ image_embeds = self._process_image_input(image_input)
369
+ inputs_embeds = merge_multimodal_embeddings(
370
+ input_ids,
371
+ inputs_embeds,
372
+ image_embeds,
373
+ placeholder_token_id=self.config.image_token_id,
374
+ )
375
+ return inputs_embeds
376
+
377
+ def forward(
378
+ self,
379
+ input_ids: Optional[torch.Tensor],
380
+ positions: torch.Tensor,
381
+ intermediate_tensors: Optional[IntermediateTensors] = None,
382
+ inputs_embeds: Optional[torch.Tensor] = None,
383
+ **kwargs,
384
+ ) -> Union[torch.Tensor, IntermediateTensors]:
385
+ if intermediate_tensors is not None:
386
+ inputs_embeds = None
387
+ elif inputs_embeds is None and kwargs.get("pixel_values") is not None:
388
+ image_input = self._parse_and_validate_image_input(**kwargs)
389
+ if image_input is None:
390
+ inputs_embeds = None
391
+ else:
392
+ assert input_ids is not None
393
+ inputs_embeds = self.get_input_embeddings_v0(
394
+ input_ids,
395
+ image_input=image_input,
396
+ )
397
+ input_ids = None
398
+
399
+ hidden_states = self.language_model(
400
+ input_ids=input_ids,
401
+ positions=positions,
402
+ intermediate_tensors=intermediate_tensors,
403
+ inputs_embeds=inputs_embeds,
404
+ )
405
+
406
+ return hidden_states
407
+
408
+ def compute_logits(
409
+ self,
410
+ hidden_states: torch.Tensor,
411
+ sampling_metadata: SamplingMetadata,
412
+ ) -> Optional[torch.Tensor]:
413
+ return self.language_model.compute_logits(hidden_states, sampling_metadata)
414
+
415
+ def sample(
416
+ self,
417
+ logits: Optional[torch.Tensor],
418
+ sampling_metadata: SamplingMetadata,
419
+ ) -> Optional[SamplerOutput]:
420
+ next_tokens = self.sampler(logits, sampling_metadata)
421
+ return next_tokens
422
+
423
+ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
424
+ loader = AutoWeightsLoader(self)
425
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
426
+
427
+
428
+ def patch_vllm_chat_placeholder():
429
+ import vllm
430
+ # return when vllm version > 0.9.1
431
+ if not (vllm.__version_tuple__[0]==0 and vllm.__version_tuple__[1] <= 9 and vllm.__version_tuple__[2] <= 1):
432
+ return
433
+ from vllm.entrypoints.chat_utils import BaseMultiModalItemTracker
434
+
435
+ ori = BaseMultiModalItemTracker._placeholder_str
436
+
437
+ def _placeholder_str(self, modality, current_count: int) -> Optional[str]:
438
+ hf_config = self._model_config.hf_config
439
+ model_type = hf_config.model_type
440
+ if modality in ("image",) and model_type in ["dots_ocr"]:
441
+ return "<|img|><|imgpad|><|endofimg|>"
442
+ return ori(self, modality, current_count)
443
+
444
+ BaseMultiModalItemTracker._placeholder_str = _placeholder_str
445
+
446
+ ModelRegistry.register_model(
447
+ "DotsOCRForCausalLM", DotsOCRForCausalLM,
448
+ )
449
+
450
+
451
+ patch_vllm_chat_placeholder()
modeling_dots_vision.py ADDED
@@ -0,0 +1,520 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ import torch.utils.checkpoint
7
+
8
+ flash_attn_available = True
9
+ npu_available = True
10
+
11
+ try:
12
+ from flash_attn import flash_attn_varlen_func
13
+ except ImportError:
14
+ flash_attn_available = False
15
+
16
+ from torch.nn import LayerNorm
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from .configuration_dots import DotsVisionConfig
19
+
20
+ try:
21
+ import torch_npu
22
+ except ImportError:
23
+ npu_available = False
24
+
25
+
26
+ def rotate_half(x):
27
+ """Rotates half the hidden dims of the input."""
28
+ x1 = x[..., : x.shape[-1] // 2]
29
+ x2 = x[..., x.shape[-1] // 2:]
30
+ return torch.cat((-x2, x1), dim=-1)
31
+
32
+
33
+ def apply_rotary_pos_emb_vision(tensor: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
34
+ orig_dtype = tensor.dtype
35
+ tensor = tensor.float()
36
+
37
+ cos = freqs.cos()
38
+ sin = freqs.sin()
39
+
40
+ cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
41
+ sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
42
+
43
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
44
+
45
+ output = output.to(orig_dtype)
46
+
47
+ return output
48
+
49
+
50
+ class VisionRotaryEmbedding(nn.Module):
51
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
52
+ super().__init__()
53
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
54
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
55
+
56
+ def forward(self, seqlen: int) -> torch.Tensor:
57
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
58
+ freqs = torch.outer(seq, self.inv_freq)
59
+ return freqs
60
+
61
+
62
+ class PatchMerger(nn.Module):
63
+ def __init__(
64
+ self,
65
+ dim: int,
66
+ context_dim: int,
67
+ spatial_merge_size: int = 2,
68
+ pre_norm="layernorm",
69
+ init_merger_std=None,
70
+ ) -> None:
71
+ super().__init__()
72
+ self.hidden_size = context_dim * (spatial_merge_size ** 2)
73
+ self.pre_norm = pre_norm
74
+ if self.pre_norm == "layernorm":
75
+ self.ln_q = LayerNorm(context_dim, eps=1e-6)
76
+ elif self.pre_norm == "rmsnorm":
77
+ self.ln_q = RMSNorm(context_dim, eps=1e-6)
78
+ else:
79
+ print("no norm in patch merger")
80
+
81
+ self.mlp = nn.Sequential(
82
+ nn.Linear(self.hidden_size, self.hidden_size),
83
+ nn.GELU(),
84
+ nn.Linear(self.hidden_size, dim),
85
+ )
86
+
87
+ if init_merger_std is not None:
88
+ nn.init.normal_(self.mlp[0].weight, mean=0.0, std=init_merger_std)
89
+ nn.init.zeros_(self.mlp[0].bias)
90
+ nn.init.normal_(self.mlp[2].weight, mean=0.0, std=init_merger_std)
91
+ nn.init.zeros_(self.mlp[2].bias)
92
+
93
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
94
+ if self.pre_norm:
95
+ x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
96
+ else:
97
+ x = self.mlp(x.view(-1, self.hidden_size))
98
+ return x
99
+
100
+
101
+ class VisionAttention(nn.Module):
102
+ def __init__(self, config, dim: int, num_heads: int = 16, bias=True) -> None:
103
+ super().__init__()
104
+ self.num_heads = num_heads
105
+ self.head_dim = dim // num_heads
106
+ self.qkv = nn.Linear(dim, dim * 3, bias=bias)
107
+ self.proj = nn.Linear(dim, dim, bias=bias)
108
+
109
+ def forward(
110
+ self,
111
+ hidden_states: torch.Tensor,
112
+ cu_seqlens: torch.Tensor,
113
+ rotary_pos_emb: torch.Tensor = None,
114
+ ) -> torch.Tensor:
115
+ seq_length = hidden_states.shape[0]
116
+
117
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
118
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
119
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
120
+
121
+ attention_mask = torch.full(
122
+ [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype
123
+ )
124
+ for i in range(1, len(cu_seqlens)):
125
+ attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = 0
126
+
127
+ q = q.transpose(0, 1)
128
+ k = k.transpose(0, 1)
129
+ v = v.transpose(0, 1)
130
+ attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
131
+ attn_weights = attn_weights + attention_mask
132
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
133
+ attn_output = torch.matmul(attn_weights, v)
134
+ attn_output = attn_output.transpose(0, 1)
135
+ attn_output = attn_output.reshape(seq_length, -1)
136
+ attn_output = self.proj(attn_output)
137
+ return attn_output
138
+
139
+
140
+ class VisionFlashAttention2(nn.Module):
141
+ def __init__(self, config, dim: int, num_heads: int = 16, bias=True) -> None:
142
+ super().__init__()
143
+ self.num_heads = num_heads
144
+ self.qkv = nn.Linear(dim, dim * 3, bias=bias)
145
+ self.proj = nn.Linear(dim, dim, bias=bias)
146
+ self.config = config
147
+ self.is_causal = config.is_causal
148
+
149
+ def forward(
150
+ self,
151
+ hidden_states: torch.Tensor,
152
+ cu_seqlens: torch.Tensor,
153
+ rotary_pos_emb: torch.Tensor = None,
154
+ ) -> torch.Tensor:
155
+ seq_length = hidden_states.shape[0]
156
+ q, k, v = (
157
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
158
+ ) # 'shd'
159
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
160
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
161
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
162
+ attn_output = flash_attn_varlen_func(
163
+ q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=self.is_causal
164
+ ).reshape(seq_length, -1)
165
+ attn_output = self.proj(attn_output)
166
+
167
+ return attn_output
168
+
169
+
170
+ class VisionAttentionV2(nn.Module):
171
+ def __init__(self, config, dim: int, num_heads: int = 16, bias=True) -> None:
172
+ super().__init__()
173
+ self.num_heads = num_heads
174
+ self.head_dim = dim // num_heads
175
+ self.qkv = nn.Linear(dim, dim * 3, bias=bias)
176
+ self.proj = nn.Linear(dim, dim, bias=bias)
177
+
178
+ def forward(
179
+ self,
180
+ hidden_states: torch.Tensor,
181
+ cu_seqlens: torch.Tensor,
182
+ rotary_pos_emb: torch.Tensor = None,
183
+ ) -> torch.Tensor:
184
+ seq_length = hidden_states.shape[0]
185
+
186
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
187
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
188
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
189
+
190
+ seqlens = torch.diff(cu_seqlens).tolist()
191
+
192
+ q_list = torch.split(q, seqlens, 0)
193
+ k_list = torch.split(k, seqlens, 0)
194
+ v_list = torch.split(v, seqlens, 0)
195
+ # eager attention 空间复杂度为 O(n^2) , n 为 b*s(batch_size * seq_len), 序列太长容易OOM, 这个实现 更具batch 切分 seq
196
+ # 减少内存需求, 计算相对 continus batching 较慢。
197
+ outputs = []
198
+ for q_i, k_i, v_i in zip(q_list, k_list, v_list):
199
+ q_i = q_i.transpose(0, 1)
200
+ k_i = k_i.transpose(0, 1)
201
+ v_i = v_i.transpose(0, 1)
202
+ out = torch.matmul(q_i, k_i.transpose(1, 2)) / math.sqrt(self.head_dim)
203
+ out = nn.functional.softmax(out, dim=-1, dtype=torch.float32).to(q.dtype)
204
+ out = torch.matmul(out, v_i)
205
+ out = out.transpose(0, 1)
206
+ outputs.append(out)
207
+
208
+ attn_output = torch.concat(outputs, dim=0)
209
+ attn_output = attn_output.reshape(seq_length, -1)
210
+ attn_output = self.proj(attn_output)
211
+ return attn_output
212
+
213
+
214
+ class VisionAscendAttention(nn.Module):
215
+ def __init__(self, config, dim: int, num_heads: int = 16, bias=True) -> None:
216
+ super().__init__()
217
+ self.num_heads = num_heads
218
+ self.head_dim = dim // num_heads
219
+ self.qkv = nn.Linear(dim, dim * 3, bias=bias)
220
+ self.proj = nn.Linear(dim, dim, bias=bias)
221
+ self.config = config
222
+
223
+ def forward(
224
+ self,
225
+ hidden_states: torch.Tensor,
226
+ cu_seqlens: torch.Tensor,
227
+ rotary_pos_emb: torch.Tensor = None,
228
+ ) -> torch.Tensor:
229
+ seq_length = hidden_states.shape[0]
230
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
231
+
232
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
233
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
234
+
235
+ attention_mask = torch.ones([1, seq_length, seq_length], device=q.device, dtype=torch.bool)
236
+ for i in range(1, len(cu_seqlens)):
237
+ attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = False
238
+
239
+ q = q.transpose(0, 1).unsqueeze(0)
240
+ k = k.transpose(0, 1).unsqueeze(0)
241
+ v = v.transpose(0, 1).unsqueeze(0)
242
+
243
+ attn_output = torch_npu.npu_prompt_flash_attention(q, k, v,
244
+ atten_mask=attention_mask,
245
+ num_heads=self.num_heads, input_layout="BNSD",
246
+ scale_value=self.head_dim ** -0.5)
247
+ attn_output = attn_output.squeeze(0).transpose(0, 1)
248
+ attn_output = attn_output.reshape(seq_length, -1)
249
+ attn_output = self.proj(attn_output)
250
+ return attn_output
251
+
252
+
253
+ class VisionSdpaAttention(nn.Module):
254
+ def __init__(self, config, dim: int, num_heads: int = 16, bias=True) -> None:
255
+ super().__init__()
256
+ self.num_heads = num_heads
257
+ self.qkv = nn.Linear(dim, dim * 3, bias=bias)
258
+ self.proj = nn.Linear(dim, dim, bias=bias)
259
+ self.config = config
260
+
261
+ def forward(
262
+ self,
263
+ hidden_states: torch.Tensor,
264
+ cu_seqlens: torch.Tensor,
265
+ rotary_pos_emb: torch.Tensor = None,
266
+ ) -> torch.Tensor:
267
+ seq_length = hidden_states.shape[0]
268
+ q, k, v = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
269
+
270
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
271
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
272
+
273
+ attention_mask = torch.zeros([1, seq_length, seq_length], device=q.device, dtype=torch.bool)
274
+ for i in range(1, len(cu_seqlens)):
275
+ attention_mask[..., cu_seqlens[i - 1]: cu_seqlens[i], cu_seqlens[i - 1]: cu_seqlens[i]] = True
276
+
277
+ # Convert q, k, v to 4D to enable : (1, num_heads, seq_length, head_dim)
278
+ q = q.transpose(0, 1).unsqueeze(0) # (1, num_heads, seq_length, head_dim)
279
+ k = k.transpose(0, 1).unsqueeze(0)
280
+ v = v.transpose(0, 1).unsqueeze(0)
281
+
282
+ # See: https://github.com/pytorch/pytorch/issues/127523
283
+ if attention_mask.stride(-1) != 1:
284
+ attention_mask = torch.empty_like(attention_mask, memory_format=torch.contiguous_format).copy_(attention_mask)
285
+
286
+ # use memory efficient backend
287
+ from torch.nn.attention import SDPBackend, sdpa_kernel
288
+ with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
289
+ attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
290
+
291
+ attn_output = attn_output.squeeze(0).transpose(0, 1) # (seq_length, num_heads, head_dim)
292
+ attn_output = attn_output.reshape(seq_length, -1)
293
+
294
+ attn_output = self.proj(attn_output)
295
+ return attn_output
296
+
297
+
298
+ DOTS_VISION_ATTENTION_CLASSES = {
299
+ "eager": VisionAttention,
300
+ "eager_v2": VisionAttentionV2, # 内存更少
301
+ "flash_attention_2": VisionFlashAttention2,
302
+ "sdpa": VisionSdpaAttention,
303
+ "ascend_fa": VisionAscendAttention, # ascend, 长序列精度下降严重。
304
+ }
305
+
306
+
307
+ class RMSNorm(nn.Module):
308
+ def __init__(self, dim: int, eps: float = 1e-6):
309
+ super().__init__()
310
+ self.weight = nn.Parameter(torch.ones(dim))
311
+ self.eps = eps
312
+
313
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
314
+ output = self._norm(x.float()).type_as(x)
315
+ return output * self.weight
316
+
317
+ def extra_repr(self) -> str:
318
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
319
+
320
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
321
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
322
+
323
+
324
+ class DotsSwiGLUFFN(nn.Module):
325
+ def __init__(self, config):
326
+ super().__init__()
327
+ hidden_features = config.intermediate_size
328
+ in_features = config.embed_dim
329
+ bias = config.use_bias
330
+
331
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
332
+ self.fc2 = nn.Linear(hidden_features, in_features, bias=bias)
333
+ self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
334
+
335
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
336
+ x = F.silu(self.fc1(x)) * self.fc3(x)
337
+ x = self.fc2(x)
338
+ return x
339
+
340
+
341
+ class DotsPatchEmbed(nn.Module):
342
+ def __init__(self, config):
343
+ super().__init__()
344
+ self.num_channels = config.num_channels
345
+ self.patch_size = config.patch_size
346
+ self.temporal_patch_size = config.temporal_patch_size
347
+ self.embed_dim = config.embed_dim
348
+ self.config = config
349
+ self.proj = nn.Conv2d(
350
+ config.num_channels,
351
+ config.embed_dim,
352
+ kernel_size=(config.patch_size, config.patch_size),
353
+ stride=(config.patch_size, config.patch_size),
354
+ )
355
+ self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
356
+
357
+ def forward(self, x: torch.Tensor, grid_thw=None) -> torch.Tensor:
358
+ x = x.view(-1, self.num_channels, self.temporal_patch_size, self.patch_size, self.patch_size)[:, :, 0]
359
+ x = self.proj(x).view(-1, self.embed_dim)
360
+ x = self.norm(x)
361
+ return x
362
+
363
+
364
+ class DotsViTPreprocessor(nn.Module):
365
+ def __init__(self, config):
366
+ super().__init__()
367
+ self.patch_h = config.patch_size
368
+ self.patch_w = config.patch_size
369
+ self.embed_dim = config.embed_dim
370
+ self.config = config
371
+ self.patchifier = DotsPatchEmbed(config)
372
+
373
+ def forward(self, x: torch.Tensor, grid_thw=None) -> torch.Tensor:
374
+ tokens = self.patchifier(x, grid_thw)
375
+ return tokens
376
+
377
+
378
+ class DotsVisionBlock(nn.Module):
379
+ def __init__(self, config, attn_implementation: str = "flash_attention_2"):
380
+ super().__init__()
381
+
382
+ if attn_implementation == "flash_attention_2" and not flash_attn_available:
383
+ # fallback to eager
384
+ attn_implementation = "eager"
385
+ print("flash attention not available! fallback to eager implementation ")
386
+
387
+ if attn_implementation == "ascend_fa" and not npu_available:
388
+ attn_implementation = "eager"
389
+ print("flash attention not available! fallback to eager implementation ")
390
+
391
+ self.attn = DOTS_VISION_ATTENTION_CLASSES[attn_implementation](
392
+ config, config.embed_dim, num_heads=config.num_attention_heads, bias=config.use_bias
393
+ )
394
+ self.norm1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
395
+ self.mlp = DotsSwiGLUFFN(config)
396
+ self.norm2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
397
+
398
+ def forward(self, hidden_states, cu_seqlens, rotary_pos_emb) -> torch.Tensor:
399
+ hidden_states = hidden_states + self.attn(
400
+ self.norm1(hidden_states), cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb
401
+ )
402
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
403
+ return hidden_states
404
+
405
+
406
+ class DotsVisionTransformer(PreTrainedModel):
407
+ def __init__(self, config: DotsVisionConfig) -> None:
408
+ super().__init__(config)
409
+ self.config = config
410
+ self.spatial_merge_size = config.spatial_merge_size
411
+
412
+ self.patch_embed = DotsViTPreprocessor(config)
413
+ self._init_weights(self.patch_embed.patchifier.proj)
414
+
415
+ head_dim = config.embed_dim // config.num_attention_heads
416
+
417
+ self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
418
+
419
+ _num_hidden_layers = config.num_hidden_layers
420
+ self.blocks = nn.ModuleList(
421
+ [DotsVisionBlock(config, config.attn_implementation) for _ in range(_num_hidden_layers)]
422
+ )
423
+
424
+ if self.config.post_norm:
425
+ self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
426
+
427
+ self.merger = PatchMerger(
428
+ dim=config.hidden_size,
429
+ context_dim=config.embed_dim,
430
+ spatial_merge_size=config.spatial_merge_size,
431
+ init_merger_std=self.config.init_merger_std,
432
+ )
433
+
434
+ self.gradient_checkpointing = False
435
+ self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
436
+
437
+ def _init_weights(self, module):
438
+ std = self.config.initializer_range
439
+ if isinstance(module, (nn.Linear, nn.Conv3d)):
440
+ module.weight.data.normal_(mean=0.0, std=std)
441
+ if module.bias is not None:
442
+ module.bias.data.zero_()
443
+ elif isinstance(module, nn.Embedding):
444
+ module.weight.data.normal_(mean=0.0, std=std)
445
+ if module.padding_idx is not None:
446
+ module.weight.data[module.padding_idx].zero_()
447
+
448
+ @property
449
+ def dtype(self) -> torch.dtype:
450
+ return self.blocks[0].mlp.fc2.weight.dtype
451
+
452
+ @property
453
+ def device(self) -> torch.device:
454
+ return self.blocks[0].mlp.fc2.weight.device
455
+
456
+ def get_pos_ids_by_grid(self, grid_thw):
457
+ pos_ids = []
458
+ for t, h, w in grid_thw:
459
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
460
+ hpos_ids = hpos_ids.reshape(
461
+ h // self.spatial_merge_size,
462
+ self.spatial_merge_size,
463
+ w // self.spatial_merge_size,
464
+ self.spatial_merge_size,
465
+ )
466
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
467
+ hpos_ids = hpos_ids.flatten()
468
+
469
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
470
+ wpos_ids = wpos_ids.reshape(
471
+ h // self.spatial_merge_size,
472
+ self.spatial_merge_size,
473
+ w // self.spatial_merge_size,
474
+ self.spatial_merge_size,
475
+ )
476
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
477
+ wpos_ids = wpos_ids.flatten()
478
+ pos_ids.append(
479
+ torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)
480
+ )
481
+
482
+ return pos_ids
483
+
484
+ def rot_pos_emb(self, grid_thw):
485
+ pos_ids = self.get_pos_ids_by_grid(grid_thw)
486
+ pos_ids = torch.cat(pos_ids, dim=0)
487
+ max_grid_size = grid_thw[:, 1:].max()
488
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
489
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
490
+ return rotary_pos_emb
491
+
492
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True) -> torch.Tensor:
493
+ if bf16:
494
+ hidden_states = hidden_states.bfloat16()
495
+ hidden_states = self.patch_embed(hidden_states, grid_thw)
496
+
497
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
498
+
499
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
500
+ dim=0,
501
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
502
+ )
503
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
504
+
505
+ for blk in self.blocks:
506
+ if self.gradient_checkpointing and self.training:
507
+ hidden_states = self._gradient_checkpointing_func(
508
+ blk.__call__,
509
+ hidden_states,
510
+ cu_seqlens,
511
+ rotary_pos_emb,
512
+ )
513
+ else:
514
+ hidden_states = blk(hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb)
515
+
516
+ if self.config.post_norm:
517
+ hidden_states = self.post_trunk_norm(hidden_states)
518
+
519
+ hidden_states = self.merger(hidden_states)
520
+ return hidden_states
preprocessor_config.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoProcessor": "configuration_dots.DotsVLProcessor"
4
+ },
5
+ "min_pixels": 3136,
6
+ "max_pixels": 11289600,
7
+ "patch_size": 14,
8
+ "temporal_patch_size": 1,
9
+ "merge_size": 2,
10
+ "image_mean": [
11
+ 0.48145466,
12
+ 0.4578275,
13
+ 0.40821073
14
+ ],
15
+ "image_std": [
16
+ 0.26862954,
17
+ 0.26130258,
18
+ 0.27577711
19
+ ],
20
+ "image_processor_type": "Qwen2VLImageProcessor",
21
+ "processor_class": "DotsVLProcessor"
22
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|object_ref_start|>",
6
+ "<|object_ref_end|>",
7
+ "<|box_start|>",
8
+ "<|box_end|>",
9
+ "<|quad_start|>",
10
+ "<|quad_end|>",
11
+ "<|vision_start|>",
12
+ "<|vision_end|>",
13
+ "<|vision_pad|>",
14
+ "<|image_pad|>",
15
+ "<|video_pad|>"
16
+ ],
17
+ "eos_token": {
18
+ "content": "<|endoftext|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ },
24
+ "pad_token": "[PAD]"
25
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,391 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "added_tokens_decoder": {
5
+ "151643": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "151644": {
14
+ "content": "<|im_start|>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "151645": {
22
+ "content": "<|im_end|>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "151646": {
30
+ "content": "<|object_ref_start|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "151647": {
38
+ "content": "<|object_ref_end|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "151648": {
46
+ "content": "<|box_start|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "151649": {
54
+ "content": "<|box_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "151650": {
62
+ "content": "<|quad_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "151651": {
70
+ "content": "<|quad_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "151652": {
78
+ "content": "<|vision_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ },
85
+ "151653": {
86
+ "content": "<|vision_end|>",
87
+ "lstrip": false,
88
+ "normalized": false,
89
+ "rstrip": false,
90
+ "single_word": false,
91
+ "special": true
92
+ },
93
+ "151654": {
94
+ "content": "<|vision_pad|>",
95
+ "lstrip": false,
96
+ "normalized": false,
97
+ "rstrip": false,
98
+ "single_word": false,
99
+ "special": true
100
+ },
101
+ "151655": {
102
+ "content": "<|image_pad|>",
103
+ "lstrip": false,
104
+ "normalized": false,
105
+ "rstrip": false,
106
+ "single_word": false,
107
+ "special": true
108
+ },
109
+ "151656": {
110
+ "content": "<|video_pad|>",
111
+ "lstrip": false,
112
+ "normalized": false,
113
+ "rstrip": false,
114
+ "single_word": false,
115
+ "special": true
116
+ },
117
+ "151657": {
118
+ "content": "<tool_call>",
119
+ "lstrip": false,
120
+ "normalized": false,
121
+ "rstrip": false,
122
+ "single_word": false,
123
+ "special": false
124
+ },
125
+ "151658": {
126
+ "content": "</tool_call>",
127
+ "lstrip": false,
128
+ "normalized": false,
129
+ "rstrip": false,
130
+ "single_word": false,
131
+ "special": false
132
+ },
133
+ "151659": {
134
+ "content": "<|fim_prefix|>",
135
+ "lstrip": false,
136
+ "normalized": false,
137
+ "rstrip": false,
138
+ "single_word": false,
139
+ "special": false
140
+ },
141
+ "151660": {
142
+ "content": "<|fim_middle|>",
143
+ "lstrip": false,
144
+ "normalized": false,
145
+ "rstrip": false,
146
+ "single_word": false,
147
+ "special": false
148
+ },
149
+ "151661": {
150
+ "content": "<|fim_suffix|>",
151
+ "lstrip": false,
152
+ "normalized": false,
153
+ "rstrip": false,
154
+ "single_word": false,
155
+ "special": false
156
+ },
157
+ "151662": {
158
+ "content": "<|fim_pad|>",
159
+ "lstrip": false,
160
+ "normalized": false,
161
+ "rstrip": false,
162
+ "single_word": false,
163
+ "special": false
164
+ },
165
+ "151663": {
166
+ "content": "<|repo_name|>",
167
+ "lstrip": false,
168
+ "normalized": false,
169
+ "rstrip": false,
170
+ "single_word": false,
171
+ "special": false
172
+ },
173
+ "151664": {
174
+ "content": "<|file_sep|>",
175
+ "lstrip": false,
176
+ "normalized": false,
177
+ "rstrip": false,
178
+ "single_word": false,
179
+ "special": false
180
+ },
181
+ "151665": {
182
+ "content": "<|imgpad|>",
183
+ "lstrip": false,
184
+ "normalized": false,
185
+ "rstrip": false,
186
+ "single_word": false,
187
+ "special": true
188
+ },
189
+ "151666": {
190
+ "content": "<|img|>",
191
+ "lstrip": false,
192
+ "normalized": false,
193
+ "rstrip": false,
194
+ "single_word": false,
195
+ "special": true
196
+ },
197
+ "151667": {
198
+ "content": "<|endofimg|>",
199
+ "lstrip": false,
200
+ "normalized": false,
201
+ "rstrip": false,
202
+ "single_word": false,
203
+ "special": true
204
+ },
205
+ "151668": {
206
+ "content": "<|systemprompt|>",
207
+ "lstrip": false,
208
+ "normalized": false,
209
+ "rstrip": false,
210
+ "single_word": false,
211
+ "special": true
212
+ },
213
+ "151669": {
214
+ "content": "<|endofsystemprompt|>",
215
+ "lstrip": false,
216
+ "normalized": false,
217
+ "rstrip": false,
218
+ "single_word": false,
219
+ "special": true
220
+ },
221
+ "151670": {
222
+ "content": "<|user|>",
223
+ "lstrip": false,
224
+ "normalized": false,
225
+ "rstrip": false,
226
+ "single_word": false,
227
+ "special": true
228
+ },
229
+ "151671": {
230
+ "content": "<|endofuser|>",
231
+ "lstrip": false,
232
+ "normalized": false,
233
+ "rstrip": false,
234
+ "single_word": false,
235
+ "special": true
236
+ },
237
+ "151672": {
238
+ "content": "<|assistant|>",
239
+ "lstrip": false,
240
+ "normalized": false,
241
+ "rstrip": false,
242
+ "single_word": false,
243
+ "special": true
244
+ },
245
+ "151673": {
246
+ "content": "<|endofassistant|>",
247
+ "lstrip": false,
248
+ "normalized": false,
249
+ "rstrip": false,
250
+ "single_word": false,
251
+ "special": true
252
+ },
253
+ "151674": {
254
+ "content": "<|ref_start|>",
255
+ "lstrip": false,
256
+ "normalized": false,
257
+ "rstrip": false,
258
+ "single_word": false,
259
+ "special": true
260
+ },
261
+ "151675": {
262
+ "content": "<|ref_end|>",
263
+ "lstrip": false,
264
+ "normalized": false,
265
+ "rstrip": false,
266
+ "single_word": false,
267
+ "special": true
268
+ },
269
+ "151676": {
270
+ "content": "[SEP]",
271
+ "lstrip": false,
272
+ "normalized": false,
273
+ "rstrip": false,
274
+ "single_word": false,
275
+ "special": true
276
+ },
277
+ "151677": {
278
+ "content": "<|pic|>",
279
+ "lstrip": false,
280
+ "normalized": false,
281
+ "rstrip": false,
282
+ "single_word": false,
283
+ "special": true
284
+ },
285
+ "151678": {
286
+ "content": "<|text|>",
287
+ "lstrip": false,
288
+ "normalized": false,
289
+ "rstrip": false,
290
+ "single_word": false,
291
+ "special": true
292
+ },
293
+ "151679": {
294
+ "content": "<|pictotext|>",
295
+ "lstrip": false,
296
+ "normalized": false,
297
+ "rstrip": false,
298
+ "single_word": false,
299
+ "special": true
300
+ },
301
+ "151680": {
302
+ "content": "[PAD]",
303
+ "lstrip": false,
304
+ "normalized": false,
305
+ "rstrip": false,
306
+ "single_word": false,
307
+ "special": true
308
+ },
309
+ "151681": {
310
+ "content": "<|slice|>",
311
+ "lstrip": false,
312
+ "normalized": false,
313
+ "rstrip": false,
314
+ "single_word": false,
315
+ "special": true
316
+ },
317
+ "151682": {
318
+ "content": "<|endofslice|>",
319
+ "lstrip": false,
320
+ "normalized": false,
321
+ "rstrip": false,
322
+ "single_word": false,
323
+ "special": true
324
+ },
325
+ "151683": {
326
+ "content": "<|imgrowend|>",
327
+ "lstrip": false,
328
+ "normalized": false,
329
+ "rstrip": false,
330
+ "single_word": false,
331
+ "special": true
332
+ },
333
+ "151684": {
334
+ "content": "<|polygon_start|>",
335
+ "lstrip": false,
336
+ "normalized": false,
337
+ "rstrip": false,
338
+ "single_word": false,
339
+ "special": true
340
+ },
341
+ "151685": {
342
+ "content": "<|polygon_end|>",
343
+ "lstrip": false,
344
+ "normalized": false,
345
+ "rstrip": false,
346
+ "single_word": false,
347
+ "special": true
348
+ },
349
+ "151686": {
350
+ "content": "<|image_gen_start|>",
351
+ "lstrip": false,
352
+ "normalized": false,
353
+ "rstrip": false,
354
+ "single_word": false,
355
+ "special": true
356
+ },
357
+ "151687": {
358
+ "content": "<|image_gen_end|>",
359
+ "lstrip": false,
360
+ "normalized": false,
361
+ "rstrip": false,
362
+ "single_word": false,
363
+ "special": true
364
+ }
365
+ },
366
+ "additional_special_tokens": [
367
+ "<|im_start|>",
368
+ "<|im_end|>",
369
+ "<|object_ref_start|>",
370
+ "<|object_ref_end|>",
371
+ "<|box_start|>",
372
+ "<|box_end|>",
373
+ "<|quad_start|>",
374
+ "<|quad_end|>",
375
+ "<|vision_start|>",
376
+ "<|vision_end|>",
377
+ "<|vision_pad|>",
378
+ "<|image_pad|>",
379
+ "<|video_pad|>"
380
+ ],
381
+ "bos_token": null,
382
+ "chat_template": "{%- for m in messages %}\n {%- if m.role == 'system' %}\n {{- '<|system|>' + m.content + '<|endofsystem|>\\n' }}\n {%- elif m.role == 'user' %}\n {{- '<|user|>' + m.content + '<|endofuser|>' }}\n {%- elif m.role == 'assistant' %}\n {{- '<|assistant|>' + m.content }}\n {%- if not loop.last %}\n {{- '<|endofassistant|>' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if messages[-1].role != 'assistant' %}\n {{- '<|assistant|>' }}\n{%- endif %}",
383
+ "clean_up_tokenization_spaces": false,
384
+ "eos_token": "<|endoftext|>",
385
+ "errors": "replace",
386
+ "model_max_length": 131072,
387
+ "pad_token": "[PAD]",
388
+ "split_special_tokens": false,
389
+ "tokenizer_class": "Qwen2Tokenizer",
390
+ "unk_token": null
391
+ }
vocab.json ADDED
The diff for this file is too large to render. See raw diff