DeepSeek-OCR Tutorial: From Messy Documents to Structured Data with Python

from transformers import AutoModel, AutoTokenizer
import torch
from PIL import Image

# Define the model repository on Hugging Face
model_name = "deepseek-ai/DeepSeek-OCR"

# Load the tokenizer
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Load the model
print("Loading model (this may download weights on first run)...")
model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
    use_safetensors=True,
    _attn_implementation="flash_attention_2" if torch.cuda.is_available() else None # Use Flash Attention if GPU is available
)

# Move model to GPU and set to bfloat16 for efficiency if CUDA is available
if torch.cuda.is_available():
    model = model.eval().cuda().to(torch.bfloat16)
else:
    model = model.eval() # CPU fallback

print("DeepSeek-OCR model loaded successfully!")

# Path to your sample image file
image_path = "sample_document.png"

try:
    image = Image.open(image_path).convert("RGB")
except FileNotFoundError:
    print(f"Error: The image file '{image_path}' was not found. Please ensure it exists.")
    exit()

# The prompt instructs the model to convert the document to Markdown
prompt = "<image>\n<|grounding|>Convert the document to markdown."

# Prepare inputs for the model
inputs = tokenizer(
    [prompt],
    [image],
    return_tensors="pt",
    padding="longest"
)

# Move inputs to GPU if available
if torch.cuda.is_available():
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

# Generate OCR output
print("Running inference on image...")
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=2048)

# Decode the output to get the Markdown text
ocr_result = tokenizer.decode(output[0], skip_special_tokens=True)

print("\n--- OCR Result (Markdown) ---")
print(ocr_result)

# Optionally, save the output to a Markdown file
output_filename = "output_image_document.md"
with open(output_filename, "w", encoding="utf-8") as f:
    f.write(ocr_result)
print(f"\nOCR result saved to {output_filename}")

import pypdfium2 as pdfium

def convert_pdf_to_images(pdf_path):
    """Converts a PDF file into a list of PIL Image objects, one per page."""
    pdf = pdfium.PdfDocument(pdf_path)
    return [page.render().to_pil() for page in pdf]

# Path to your sample PDF file
pdf_path = "sample_document.pdf"
all_pages_markdown = ""

try:
    pdf_images = convert_pdf_to_images(pdf_path)
except FileNotFoundError:
    print(f"Error: The PDF file '{pdf_path}' was not found. Please ensure it exists.")
    exit()

print(f"Processing {len(pdf_images)} pages from PDF...")
for i, page_image in enumerate(pdf_images):
    print(f"  Processing page {i+1}/{len(pdf_images)}...")
    
    # Use the same inference logic as for single images
    inputs = tokenizer([prompt], [page_image], return_tensors="pt", padding="longest")
    if torch.cuda.is_available():
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=2048)
    
    page_markdown = tokenizer.decode(output[0], skip_special_tokens=True)
    all_pages_markdown += f"\n\n--- Page {i+1} ---\n\n" + page_markdown # Add a separator between pages

# Save the combined Markdown output for the entire PDF
output_pdf_filename = "output_pdf_document.md"
with open(output_pdf_filename, "w", encoding="utf-8") as f:
    f.write(all_pages_markdown)
print(f"\nPDF processing complete. Combined Markdown saved to {output_pdf_filename}")