add transformers usage of PaddleOCR-VL-0.9B

#32
Files changed (2) hide show
  1. README.md +52 -0
  2. chat_template.jinja +27 -3
README.md CHANGED
@@ -140,6 +140,58 @@ for res in output:
140
  ```
141
 
142
  **For more usage details and parameter explanations, see the [documentation](https://www.paddleocr.ai/latest/en/version3.x/pipeline_usage/PaddleOCR-VL.html).**
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
143
  ## Performance
144
 
145
  ### Page-Level Document Parsing
 
140
  ```
141
 
142
  **For more usage details and parameter explanations, see the [documentation](https://www.paddleocr.ai/latest/en/version3.x/pipeline_usage/PaddleOCR-VL.html).**
143
+
144
+ ## PaddleOCR-VL-0.9B Usage with transformers
145
+
146
+
147
+ Currently, we support inference using the PaddleOCR-VL-0.9B model with the `transformers` library, which can recognize texts, formulas, tables, and chart elements. In the future, we plan to support full document parsing inference with `transformers`. Below is a simple script we provide to support inference using the PaddleOCR-VL-0.9B model with `transformers`. We currently recommend using the official method for inference, which is faster and can support page-level document parsing.
148
+
149
+
150
+ ```python
151
+ from PIL import Image
152
+ import torch
153
+ from transformers import AutoModelForCausalLM, AutoProcessor
154
+
155
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
156
+
157
+ CHOSEN_TASK = "ocr" # Options: 'ocr' | 'table' | 'chart' | 'formula'
158
+ PROMPTS = {
159
+ "ocr": "OCR:",
160
+ "table": "Table Recognition:",
161
+ "formula": "Formula Recognition:",
162
+ "chart": "Chart Recognition:",
163
+ }
164
+
165
+ model_path = "PaddlePaddle/PaddleOCR-VL"
166
+ image_path = "test.png"
167
+ image = Image.open(image_path).convert("RGB")
168
+
169
+ model = AutoModelForCausalLM.from_pretrained(
170
+ model_path, trust_remote_code=True, torch_dtype=torch.bfloat16
171
+ ).to(DEVICE).eval()
172
+ processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
173
+
174
+ messages = [
175
+ {"role": "user",
176
+ "content": [
177
+ {"type": "image", "image": image},
178
+ {"type": "text", "text": PROMPTS[CHOSEN_TASK]},
179
+ ]
180
+ }
181
+ ]
182
+ inputs = processor.apply_chat_template(
183
+ messages,
184
+ tokenize=True,
185
+ add_generation_prompt=True,
186
+ return_dict=True,
187
+ return_tensors="pt"
188
+ ).to(DEVICE)
189
+
190
+ outputs = model.generate(**inputs, max_new_tokens=1024)
191
+ outputs = processor.batch_decode(outputs, skip_special_tokens=True)[0]
192
+ print(outputs)
193
+ ```
194
+
195
  ## Performance
196
 
197
  ### Page-Level Document Parsing
chat_template.jinja CHANGED
@@ -7,14 +7,38 @@
7
  {%- if not sep_token is defined -%}
8
  {%- set sep_token = "<|end_of_sentence|>" -%}
9
  {%- endif -%}
 
 
 
10
  {{- cls_token -}}
11
  {%- for message in messages -%}
12
  {%- if message["role"] == "user" -%}
13
- {{- "User: <|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>" + message["content"] + "\n" -}}
 
 
 
 
 
 
 
 
 
 
 
14
  {%- elif message["role"] == "assistant" -%}
15
- {{- "Assistant: " + message["content"] + sep_token -}}
 
 
 
 
 
 
16
  {%- elif message["role"] == "system" -%}
17
- {{- message["content"] -}}
 
 
 
 
18
  {%- endif -%}
19
  {%- endfor -%}
20
  {%- if add_generation_prompt -%}
 
7
  {%- if not sep_token is defined -%}
8
  {%- set sep_token = "<|end_of_sentence|>" -%}
9
  {%- endif -%}
10
+ {%- if not image_token is defined -%}
11
+ {%- set image_token = "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>" -%}
12
+ {%- endif -%}
13
  {{- cls_token -}}
14
  {%- for message in messages -%}
15
  {%- if message["role"] == "user" -%}
16
+ {{- "User: " -}}
17
+ {%- for content in message["content"] -%}
18
+ {%- if content["type"] == "image" -%}
19
+ {{ image_token }}
20
+ {%- endif -%}
21
+ {%- endfor -%}
22
+ {%- for content in message["content"] -%}
23
+ {%- if content["type"] == "text" -%}
24
+ {{ content["text"] }}
25
+ {%- endif -%}
26
+ {%- endfor -%}
27
+ {{ "\n" -}}
28
  {%- elif message["role"] == "assistant" -%}
29
+ {{- "Assistant: " -}}
30
+ {%- for content in message["content"] -%}
31
+ {%- if content["type"] == "text" -%}
32
+ {{ content["text"] + "\n" }}
33
+ {%- endif -%}
34
+ {%- endfor -%}
35
+ {{ sep_token -}}
36
  {%- elif message["role"] == "system" -%}
37
+ {%- for content in message["content"] -%}
38
+ {%- if content["type"] == "text" -%}
39
+ {{ content["text"] + "\n" }}
40
+ {%- endif -%}
41
+ {%- endfor -%}
42
  {%- endif -%}
43
  {%- endfor -%}
44
  {%- if add_generation_prompt -%}