Views
No views yet

Dots.OCR-Latest-BF16 is an optimized and updated vision-language OCR model variant of the original Dots.OCR. This open-source model is designed to extract text from images and scanned documents, including handwritten and printed content. It can output results as plain text or Markdown, preserving document layout elements such as headings, tables, and lists. This model uses a powerful multimodal backbone (3B VLM) to enhance reading comprehension and layout understanding, handling cursive handwriting and complex document structures effectively.
transformers version without compatibility issues, ensuring optimized performance.transformers: 4.57.1
torch: 2.6.0+cu124
cuda: 12.4
device: NVIDIA H200 MIG 3g.71gb
attn_implementation= "flash_attention_2"gradio
numpy
torch
torchvision
transformers==4.57.1
accelerate
matplotlib
flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl1import os
2import sys
3import random
4import uuid
5import json
6import time
7from threading import Thread
8from typing import Iterable
9from huggingface_hub import snapshot_download
10
11import gradio as gr
12import torch
13import numpy as np
14from PIL import Image
15import cv2
16
17from transformers import (
18 AutoModelForCausalLM,
19 AutoProcessor,
20 TextIteratorStreamer,
21)
22
23from transformers.image_utils import load_image
24
25css = """
26#main-title h1 {
27 font-size: 2.3em !important;
28}
29#output-title h2 {
30 font-size: 2.1em !important;
31}
32"""
33
34MAX_MAX_NEW_TOKENS = 4096
35DEFAULT_MAX_NEW_TOKENS = 2048
36MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
37
38device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
39
40print("--- System Information ---")
41print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
42print("torch.__version__ =", torch.__version__)
43print("torch.version.cuda =", torch.version.cuda)
44print("CUDA available:", torch.cuda.is_available())
45print("CUDA device count:", torch.cuda.device_count())
46if torch.cuda.is_available():
47 print("Current device:", torch.cuda.current_device())
48 print("Device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
49print("Using device:", device)
50print("--------------------------")
51
52print("Loading Dots.OCR model...")
53MODEL_PATH_D = "prithivMLmods/Dots.OCR-Latest-BF16"
54processor = AutoProcessor.from_pretrained(MODEL_PATH_D, trust_remote_code=True)
55model = AutoModelForCausalLM.from_pretrained(
56 MODEL_PATH_D,
57 attn_implementation="flash_attention_2",
58 torch_dtype=torch.bfloat16,
59 device_map="auto",
60 trust_remote_code=True
61).eval()
62print("Dots.OCR model loaded successfully.")
63
64def generate_image(text: str, image: Image.Image,
65 max_new_tokens: int, temperature: float, top_p: float,
66 top_k: int, repetition_penalty: float):
67 """
68 Generates responses using the Dots.OCR model for image input.
69 Yields raw text and Markdown-formatted text.
70 """
71 if image is None:
72 yield "Please upload an image.", "Please upload an image."
73 return
74
75 messages = [{
76 "role": "user",
77 "content": [
78 {"type": "image"},
79 {"type": "text", "text": text},
80 ]
81 }]
82 prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
83
84 inputs = processor(
85 text=[prompt_full],
86 images=[image],
87 return_tensors="pt",
88 padding=True).to(device)
89
90 streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
91 generation_kwargs = {
92 **inputs,
93 "streamer": streamer,
94 "max_new_tokens": max_new_tokens,
95 "do_sample": True,
96 "temperature": temperature,
97 "top_p": top_p,
98 "top_k": top_k,
99 "repetition_penalty": repetition_penalty,
100 }
101 thread = Thread(target=model.generate, kwargs=generation_kwargs)
102 thread.start()
103 buffer = ""
104 for new_text in streamer:
105 buffer += new_text
106 # Clean up potential end-of-sequence tokens from the buffer
107 buffer = buffer.replace("<|im_end|>", "")
108 time.sleep(0.01)
109 yield buffer, buffer
110
111with gr.Blocks(css=css) as demo:
112 gr.Markdown("# **Dots.OCR**", elem_id="main-title")
113 with gr.Row():
114 with gr.Column(scale=2):
115 image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
116 image_upload = gr.Image(type="pil", label="Upload Image", height=290)
117
118 image_submit = gr.Button("Submit", variant="primary")
119
120 with gr.Accordion("Advanced options", open=False):
121 max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
122 temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.7)
123 top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
124 top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
125 repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.1)
126
127 with gr.Column(scale=3):
128 gr.Markdown("## Output", elem_id="output-title")
129 output = gr.Textbox(label="Raw Output Stream", interactive=False, lines=15, show_copy_button=True)
130 with gr.Accordion("(Result.md)", open=False):
131 markdown_output = gr.Markdown(label="(Result.Md)")
132
133 image_submit.click(
134 fn=generate_image,
135 inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
136 outputs=[output, markdown_output]
137 )
138
139if __name__ == "__main__":
140 demo.queue(max_size=50).launch(ssr_mode=False, show_error=True)| Resource Type | Description | Link |
|---|---|---|
| Original Model Card | Official release of Dots.OCR by rednote-hilab | rednote-hilab/dots.ocr |
| Test Model (StrangerZone HF) | Community test deployment (experimental) | strangervisionhf/dots.ocr-base-fix |
| Standard Model Card | Optimized version supporting Transformers v4.57.1 (BF16 precision) | prithivMLmods/Dots.OCR-Latest-BF16 |
| Demo Space | Interactive demo hosted on Hugging Face Spaces | Multimodal-OCR3 Demo |