DeepSeek-OCR-2

https://huggingface.co/deepseek-ai/DeepSeek-OCR-2
Admin
DeepSeek-OCR-2 is a vision-language model that performs optical character recognition and document...
Tokens:73,909
Snippets:351
Trust Score:10
Update:8 months ago
Show doc for...
Tokens:
Raw
### DeepSeek-OCR-2 Prompt Examples

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/16dca681b59d4b6a77657c827051547bb69faa5c

Provides various prompt examples for different use cases with the DeepSeek-OCR-2 model. These examples illustrate how to guide the model for tasks like document conversion, OCR, figure parsing, and detailed image description.

```python
# document: <image>\n<|grounding|>Convert the document to markdown.
# other image: <image>\n<|grounding|>OCR this image.
# without layouts: <image>\nFree OCR.
# figures in document: <image>\nParse the figure.
# general: <image>\nDescribe this image in detail.
# rec: <image>\nLocate <|ref|>xxxx<|/ref|> in the image.
```

--------------------------------

### Prompt Examples for DeepSeek-OCR-2

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/db8cee3a6b6a12ff7237046281a0523579ee74fa

This section provides various prompt examples for interacting with the DeepSeek-OCR-2 model. These prompts illustrate different use cases, including document conversion to markdown, general OCR, layout analysis, and describing or locating elements within an image.

```python
# document: <image>
<|grounding|>Convert the document to markdown.
# other image: <image>
<|grounding|>OCR this image.
# without layouts: <image>
Free OCR.
# figures in document: <image>
Parse the figure.
# general: <image>
Describe this image in detail.
# rec: <image>
Locate <|ref|>xxxx<|/ref|> in the image.
```

--------------------------------

### Install DeepSeek-OCR-2 Dependencies

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Installs the required Python packages for DeepSeek-OCR-2, including PyTorch, Transformers, and Flash Attention. The `--no-build-isolation` flag is used for Flash Attention installation.

```bash
pip install torch==2.6.0
pip install transformers==4.46.3
pip install tokenizers==0.20.3
pip install einops
pip install addict
pip install easydict
pip install flash-attn==2.7.3 --no-build-isolation
```

--------------------------------

### DeepseekV2ForCausalLM Initialization and Core Methods (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet shows the initialization of the DeepseekV2ForCausalLM class, which inherits from DeepseekV2PreTrainedModel. It also includes methods for getting and setting input embeddings, output embeddings, and the decoder. These methods are essential for model manipulation and configuration.

```python
class DeepseekV2ForCausalLM(DeepseekV2PreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.model = DeepseekV2Model(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model
```

--------------------------------

### Perform OCR with DeepSeek-OCR 2

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/discussions/10/files

This example shows how to use the loaded DeepSeek-OCR 2 model to perform OCR on an image. It defines a prompt for converting document content to markdown and specifies input and output paths. This requires an image file and a directory to save the output.

```python
# prompt = "<image> \nConvert the document to markdown. "
# image_file = 'your_image.jpg'
# output_path = 'your/output/dir'

# Note: The actual inference call would follow, utilizing the 'model' and 'tokenizer' objects.
```

--------------------------------

### Install DeepSeek-OCR-2 Dependencies

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/blame/main/README

Installs the necessary Python libraries for using the DeepSeek-OCR-2 model with Hugging Face Transformers. This includes torch, transformers, tokenizers, flash-attn, and other dependencies. Ensure CUDA is available for GPU acceleration.

```bash
pip install torch==2.6.0
transformers==4.46.3
tokenizers==0.20.3
einops
addict 
easydict
pip install flash-attn==2.7.3 --no-build-isolation
```

--------------------------------

### Install DeepSeek-OCR-2 Dependencies with Transformers

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/blob/main/README

Installs the necessary libraries for using the DeepSeek-OCR-2 model with Hugging Face Transformers. This includes PyTorch, Transformers, tokenizers, and flash-attn for optimized attention mechanisms. Ensure you have a compatible Python version and CUDA installed.

```bash
pip install flash-attn==2.7.3 --no-build-isolation
```

--------------------------------

### Initialize and Use DeepSeek Conversation Template

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet demonstrates how to get the 'deepseek' conversation template, append messages from different roles (user and assistant), and then print the generated prompt. It requires the 'get_conv_template' function, likely from a library managing AI conversation structures.

```python
if __name__ == "__main__":
  print("deepseek template:")
  conv = get_conv_template("deepseek")
  conv.append_message(conv.roles[0], "Hello!")
  conv.append_message(conv.roles[1], "Hi! This is Tony.")
  conv.append_message(conv.roles[0], "Who are you?")
  conv.append_message(conv.roles[1], "I am a helpful assistant.")
  conv.append_message(conv.roles[0], "How are you?")
  conv.append_message(conv.roles[1], None)
  print(conv.get_prompt())
```

--------------------------------

### Install DeepSeek-OCR-2 Dependencies

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/blob/main/README_code=true

This snippet lists the necessary Python packages and specific versions required for running the DeepSeek-OCR-2 model with Hugging Face Transformers. It includes core libraries like PyTorch and Transformers, along with specialized packages for attention mechanisms.

```text
torch==2.6.0
transformers==4.46.3
tokenizers==0.20.3
einops
addict
easydict
pip install flash-attn==2.7.3 --no-build-isolation
```

--------------------------------

### Load DeepSeek-OCR 2 Model with Transformers

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/discussions/10/files

This snippet demonstrates how to load the DeepSeek-OCR 2 model and its tokenizer using the Hugging Face Transformers library. It specifies the model name and configures it for efficient inference with Flash Attention 2 and bfloat16 precision. Ensure you have the 'transformers', 'torch', and 'accelerate' libraries installed.

```python
from transformers import AutoTokenizer, AutoModel

model_name = "deepseek-ai/deepseek-ocr-2-turbo"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, _attn_implementation='flash_attention_2', trust_remote_code=True, use_safetensors=True)
model = model.eval().cuda().to(torch.bfloat16)
```

--------------------------------

### DeepSeek-OCR-2 Main Prompts

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/blob/main/README_code=true

This snippet provides example prompts for using the DeepSeek-OCR-2 model for different OCR tasks. It includes prompts for document conversion to markdown and a general free OCR task.

```python
# document: <image>\n<|grounding|>Convert the document to markdown.
# without layouts: <image>\nFree OCR.
```

--------------------------------

### MoEGate Module for Mixture-of-Experts (PyTorch)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Defines the MoEGate module, responsible for routing tokens to experts in a Mixture-of-Experts (MoE) setup. It incorporates parameters for top-k routing, expert scaling, scoring functions, and auxiliary loss calculations, supporting various routing strategies.

```python
class MoEGate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.routed_scaling_factor = config.routed_scaling_factor
        self.scoring_func = config.scoring_func
        self.alpha = config.aux_loss_alpha
        self.seq_aux = config.seq_aux
        self.topk_method = config.topk_method
        self.n_group = config.n_group
        self.topk_group = config.topk_group
```

--------------------------------

### Initialize DeepSeek OCR 2 Model with PyTorch

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet demonstrates the initialization of the DeepSeek OCR 2 model using PyTorch. It imports necessary libraries such as torch, transformers, and specific components for attention mechanisms and model configuration. The code is designed to be compatible with DeepSeekV2 and DeepSeekV3 architectures.

```python
import math
import warnings
from typing import List, Optional, Tuple, Union
import numpy as np

import torch
import torch.nn.functional as F
import torch.utils.checkpoint
import torch.distributed as dist
from einops import repeat
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
from transformers.models.llama.modeling_llama import (
    LlamaAttention,
    LlamaFlashAttention2
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    SequenceClassifierOutputWithPast,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.pytorch_utils import (
    ALL_LAYERNORM_LAYERS,
    is_torch_greater_or_equal_than_1_13,
)
from transformers.utils import (
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    is_flash_attn_2_available,
    is_flash_attn_greater_or_equal_2_10,
    logging,
    replace_return_docstrings,
)
from transformers.utils.import_utils import is_torch_fx_available

from .configuration_deepseek_v2 import DeepseekV2Config

if is_flash_attn_2_available():
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa

# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
# It means that the function will not be traced through and simply appear as a node in the graph.
if is_torch_fx_available():
    if not is_torch_greater_or_equal_than_1_13:
        import torch.fx

_prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)

logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "DeepseekV2Config"
```

--------------------------------

### Reset Conversation (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Resets the conversation by clearing all stored messages. This function is used to start a new dialogue or clear the current context.

```python
self.messages = []
```

--------------------------------

### Initialize and Use DeepSeekV2 Conversation Template

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet shows the process of obtaining the 'deepseekv2' conversation template, adding messages from defined roles, and finally outputting the complete prompt. Similar to the 'deepseek' template, it relies on a function to retrieve conversation templates.

```python
print("deepseekv2 template:")
conv = get_conv_template("deepseekv2")
conv.append_message(conv.roles[0], "Hello!")
conv.append_message(conv.roles[1], "Hi! This is Tony.")
conv.append_message(conv.roles[0], "Who are you?")
conv.append_message(conv.roles[1], "I am a helpful assistant.")
conv.append_message(conv.roles[0], "How are you?")
conv.append_message(conv.roles[1], None)
print(conv.get_prompt())
```

--------------------------------

### Get Conversation Template (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Retrieves a copy of a registered conversation template by its name. This ensures that modifications to the retrieved template do not affect the original registered template.

```python
def get_conv_template(name: str) -> Conversation:
    """Get a conversation template."""
    return conv_templates[name].copy()
```

--------------------------------

### DeepseekV2PreTrainedModel Initialization and Configuration

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Defines the DeepseekV2PreTrainedModel class, inheriting from PreTrainedModel. It sets configuration class, base model prefix, and flags for gradient checkpointing, module splitting, and Flash Attention 2 support. It also defines which keys should be skipped during device placement.

```python
class DeepseekV2PreTrainedModel(PreTrainedModel):
    config_class = DeepseekV2Config
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["DeepseekV2DecoderLayer"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True
    _supports_cache_class = True
```

--------------------------------

### Add BOS Token to Input Sequence

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Prepends a Beginning-Of-Sentence (BOS) token ID to the tokenized string and its corresponding mask. This is a standard practice in sequence modeling to indicate the start of a sequence.

```python
bos_id = 0
tokenized_str = [bos_id] + tokenized_str
images_seq_mask = [False] + images_seq_mask
```

--------------------------------

### Input Embedding Management

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Provides methods to get and set the input embedding layer of the DeepseekV2Model. This allows for direct access and modification of how input tokens are converted into dense vector representations.

```python
def get_input_embeddings(self):
    return self.embed_tokens

def set_input_embeddings(self, value):
    self.embed_tokens = value
```

--------------------------------

### Initialize DeepseekV2Config in Python

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Demonstrates how to initialize a DeepseekV2Config object using default parameters. This configuration is essential for setting up the Deepseek-V2 model architecture.

```python
from transformers import DeepseekV2Model, DeepseekV2Config

# Initializing a Deepseek-V2 style configuration
configuration = DeepseekV2Config()
```

--------------------------------

### Custom Qwen2 Decoder Initialization and Configuration (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/18bdb3c4c8ccd6fe508f9fdedcd6438b01fcd3c5

Initializes a custom Qwen2 decoder by loading the Qwen2Model and Qwen2Config from the transformers library. It sets up the model configuration based on provided parameters and includes a check to prevent the use of 'flash_attention_2'.

```python
class CustomQwen2Decoder(nn.Module):
    """
    Qwen2 visual encoder
    non-causal attention + causal attention
    token_type_ids ：0=non-causal, 1=causal
    """

    def __init__(self,
            decoder_layer: int = 24,
            max_position_embeddings: int = 131072,
            hidden_dimension: int = 896,
            num_attention_heads: int = 14,
            num_key_value_heads: int = 2,
            intermediate_size: int = 4864,
            vocab_size: int = 151936,
            attn_implementation: str = "sdpa", # ⭐
            rms_norm_eps: float = 1e-06,
            rope_theta: float = 1000000.0,
            attention_dropout: float = 0.0,
            hidden_act: str = "silu",
            initializer_range: float = 0.02,
            ):
        super().__init__()

        # attn_implementation check
        if attn_implementation == "flash_attention_2":
            raise ValueError(
                "CustomQwen2Decoder do not support flash_attention_2，"
                "new attention mask needs 'sdpa' or 'eager'"
            )

        # load
        Qwen2Model = getattr(transformers.models.qwen2.modeling_qwen2, 'Qwen2Model')
        Qwen2Config = getattr(transformers, 'Qwen2Config')

        # config
        config = Qwen2Config(
            hidden_size=hidden_dimension,
            num_hidden_layers=decoder_layer,
            num_attention_heads=num_attention_heads,
            num_key_value_heads=num_key_value_heads,
            intermediate_size=intermediate_size,
            max_position_embeddings=max_position_embeddings,
            vocab_size=vocab_size,
            rms_norm_eps=rms_norm_eps,
            rope_theta=rope_theta,
            attention_dropout=attention_dropout,
            hidden_act=hidden_act,
            initializer_range=initializer_range,
            _attn_implementation=attn_implementation, # ⭐
            )

        self.model = self._create_custom_model(Qwen2Model, config)

        del self.model.embed_tokens
```

--------------------------------

### DeepseekOCR2ForCausalLM Initialization

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet details the initialization of the DeepseekOCR2ForCausalLM class. It sets up the configuration, instantiates the DeepseekOCR2Model, and defines the language modeling head (lm_head) for predicting vocabulary tokens.

```python
class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
    config_class = DeepseekOCR2Config
    # supports_gradient_checkpointing = True

    def __init__(self, config):
        super(DeepseekV2ForCausalLM, self).__init__(config)
        self.model = DeepseekOCR2Model(config)

        self.vocab_size = config.vocab_size

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()
```

--------------------------------

### DeepseekV2Model Initialization

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Initializes the DeepseekV2Model, setting up embedding layers, decoder layers, and normalization. It configures the model based on the provided DeepseekV2Config, including embedding dimensions, number of layers, and attention implementation.

```python
class DeepseekV2Model(DeepseekV2PreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV2DecoderLayer`]
    
    Args:
    config: DeepseekV2Config
    """
    
    def __init__(self, config: DeepseekV2Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        
        self.embed_tokens = nn.Embedding(
            config.vocab_size, config.hidden_size, self.padding_idx
        )
        self.layers = nn.ModuleList(
            [
                DeepseekV2DecoderLayer(config, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        # print(config._attn_implementation)
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.norm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()
```

--------------------------------

### DeepseekV2Attention Module Initialization in PyTorch

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet shows the initialization of the `DeepseekV2Attention` module, a custom multi-headed attention layer. It configures various parameters like hidden size, number of heads, RoPE (Rotary Positional Embedding) settings, and LoRA (Low-Rank Adaptation) ranks for query and key-value projections. It utilizes `nn.Linear` for projections and `DeepseekV2RMSNorm` for normalization.

```python
class DeepseekV2Attention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""
    
    def __init__(self, config: DeepseekV2Config, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        if layer_idx is None:
            logger.warning_once(
                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx""
                "when creating this class."
            )
            
        self.attention_dropout = config.attention_dropout
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta
        self.q_lora_rank = config.q_lora_rank
        self.qk_rope_head_dim = config.qk_rope_head_dim
        self.kv_lora_rank = config.kv_lora_rank
        self.v_head_dim = config.v_head_dim
        self.qk_nope_head_dim = config.qk_nope_head_dim
        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
        
        self.is_causal = True
        
        if self.q_lora_rank is None:
            self.q_proj = nn.Linear(
                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
            )
        else:
            self.q_a_proj = nn.Linear(
                self.hidden_size, config.q_lora_rank, bias=config.attention_bias
            )
            self.q_a_layernorm = DeepseekV2RMSNorm(config.q_lora_rank)
```

--------------------------------

### DeepSeek-OCR Model Initialization Parameters

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Initializes the DeepSeek-OCR model with various configuration parameters. These include settings for attention heads, vocabulary size, hidden layer dimensions, and specialized parameters for Mixture of Experts (MoE) layers. It also handles backward compatibility for `num_key_value_heads`.

```python
def __init__(
    self,
    vocab_size,
    hidden_size,
    intermediate_size,
    moe_intermediate_size,
    num_hidden_layers,
    num_attention_heads,
    n_shared_experts,
    n_routed_experts,
    ep_size,
    routed_scaling_factor,
    kv_lora_rank,
    q_lora_rank,
    qk_rope_head_dim = 64,
    v_head_dim = 128,
    qk_nope_head_dim = 128,
    topk_method = 'gready',
    n_group = None,
    topk_group = None,
    num_experts_per_tok = None,
    moe_layer_freq = 1,
    first_k_dense_replace = 0,
    norm_topk_prob = False,
    scoring_func = 'softmax',
    aux_loss_alpha = 0.001,
    seq_aux = True,
    hidden_act="silu",
    max_position_embeddings=2048,
    initializer_range=0.02,
    rms_norm_eps=1e-6,
    use_cache=True,
    pad_token_id=None,
    bos_token_id=100000,
    eos_token_id=100001,
    pretraining_tp=1,
    tie_word_embeddings=False,
    rope_theta=10000.0,
    rope_scaling=None,
    attention_bias=False,
    attention_dropout=0.0,
    use_mla=True,
    num_key_value_heads = None,
    **kwargs,
):
    self.vocab_size = vocab_size
    self.max_position_embeddings = max_position_embeddings
    self.hidden_size = hidden_size
    self.intermediate_size = intermediate_size
    self.moe_intermediate_size = moe_intermediate_size
    self.num_hidden_layers = num_hidden_layers
    self.num_attention_heads = num_attention_heads
    self.n_shared_experts = n_shared_experts
    self.n_routed_experts = n_routed_experts
    self.ep_size = ep_size
    self.routed_scaling_factor = routed_scaling_factor
    self.kv_lora_rank = kv_lora_rank
    self.q_lora_rank = q_lora_rank
    self.qk_rope_head_dim = qk_rope_head_dim
    self.v_head_dim = v_head_dim
    self.qk_nope_head_dim = qk_nope_head_dim
    self.topk_method = topk_method
    self.n_group = n_group
    self.topk_group = topk_group
    self.num_experts_per_tok = num_experts_per_tok
    self.moe_layer_freq = moe_layer_freq
    self.first_k_dense_replace = first_k_dense_replace
    self.norm_topk_prob = norm_topk_prob
    self.scoring_func = scoring_func
    self.aux_loss_alpha = aux_loss_alpha
    self.seq_aux = seq_aux
    # for backward compatibility
    if num_key_value_heads is None:
        num_key_value_heads = num_attention_heads
    self.num_key_value_heads = num_key_value_heads
    self.hidden_act = hidden_act
    self.initializer_range = initializer_range
    self.rms_norm_eps = float(rms_norm_eps)
    self.pretraining_tp = pretraining_tp
    self.use_cache = use_cache
    self.rope_theta = rope_theta
    self.rope_scaling = rope_scaling
    self.attention_bias = attention_bias
    self.attention_dropout = attention_dropout
    self.use_mla = use_mla

    super().__init__(
        pad_token_id=pad_token_id,
        bos_token_id=bos_token_id,
        eos_token_id=eos_token_id,
        tie_word_embeddings=tie_word_embeddings,
        **kwargs,
    )

```

--------------------------------

### Initialize and Load Image Encoder in Python

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet initializes an ImageEncoderViT model with specified parameters and loads a pre-trained checkpoint. It handles checkpoint loading, including potential key transformations for matching weights, and returns the configured image encoder. Dependencies include the torch library.

```python
image_embedding_size = image_size // vit_patch_size
image_encoder=ImageEncoderViT(
    depth=encoder_depth,
    embed_dim=encoder_embed_dim,
    img_size=image_size,
    mlp_ratio=4,
    norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
    num_heads=encoder_num_heads,
    patch_size=vit_patch_size,
    qkv_bias=True,
    use_rel_pos=True,
    global_attn_indexes=encoder_global_attn_indexes,
    window_size=14,
    out_chans=prompt_embed_dim,
    )
image_encoder.eval()
if checkpoint is not None:
    state_dict = torch.load(checkpoint)
    image_encoder.load_state_dict({k[30:]: v for k, v in state_dict.items() if 'vision_tower_high' in k}, strict=True)
print(checkpoint)
return image_encoder
```

--------------------------------

### Get Relative Positional Embeddings (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/db8cee3a6b6a12ff7237046281a0523579ee74fa

Calculates relative positional embeddings based on query and key sizes. It handles interpolation if the provided relative position embeddings do not match the required size. This function is crucial for attention mechanisms that incorporate relative spatial information.

```python
def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
    """
    Get relative positional embeddings according to the relative positions of
    query and key sizes.
    Args:
    q_size (int): size of query q.
    k_size (int): size of key k.
    rel_pos (Tensor): relative position embeddings (L, C).

    Returns:
    Extracted positional embeddings according to relative positions.
    """
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
    # Interpolate rel pos.
    dtype = rel_pos.dtype
    rel_pos = rel_pos.to(torch.float32)
    rel_pos_resized = F.interpolate(
    rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
    size=max_rel_dist,
    mode="linear",
    ).to(dtype)
    rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
    else:
    rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]
```

--------------------------------

### DeepseekV2MoE Module Initialization

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Initializes the DeepseekV2MoE module, configuring shared experts and handling distributed settings. It sets up the experts (MLP layers) and the gating mechanism based on the provided configuration. If distributed training is enabled (ep_size > 1), it partitions experts across ranks.

```python
class DeepseekV2MoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts_per_tok = config.num_experts_per_tok

        if hasattr(config, "ep_size") and config.ep_size > 1:
            assert config.ep_size == dist.get_world_size()
            self.ep_size = config.ep_size
            self.experts_per_rank = config.n_routed_experts // config.ep_size
            self.ep_rank = dist.get_rank()
            self.experts = nn.ModuleList(
                [
                    (
                        DeepseekV2MLP(
                            config, intermediate_size=config.moe_intermediate_size
                        )
                        if i >= self.ep_rank * self.experts_per_rank
                        and i < (self.ep_rank + 1) * self.experts_per_rank
                        else None
                    )
                    for i in range(config.n_routed_experts)
                ]
            )
        else:
            self.ep_size = 1
            self.experts_per_rank = config.n_routed_experts
            self.ep_rank = 0
            self.experts = nn.ModuleList(
                [
                    DeepseekV2MLP(
                        config, intermediate_size=config.moe_intermediate_size
                    )
                    for i in range(config.n_routed_experts)
                ]
            )
        self.gate = MoEGate(config)
        if config.n_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = DeepseekV2MLP(
                config=config, intermediate_size=intermediate_size
            )
```

--------------------------------

### Conversation Template Management in Python

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This Python code defines a Conversation class to manage prompt templates and conversation history. It supports various separator styles like DeepSeek, DeepSeekV2, PLAIN, and ALIGNMENT, and includes methods for getting the formatted prompt, setting system messages, and appending/updating conversation turns.

```python
import dataclasses
from enum import IntEnum, auto
from typing import Any, Dict, List

class SeparatorStyle(IntEnum):
    """Separator styles."""
    DeepSeek = auto()
    DeepSeekV2 = auto()
    PLAIN = auto()
    ALIGNMENT = auto()

@dataclasses.dataclass
class Conversation:
    """A class that manages prompt templates and keeps all conversation history."""
    name: str
    system_template: str = "{system_message}"
    system_message: str = ""
    roles: List[str] = (("USER", "ASSISTANT"),)
    messages: List[List[str]] = ()
    offset: int = 0
    sep_style: SeparatorStyle = SeparatorStyle.DeepSeek
    sep: str = "\n"
    sep2: str = None
    stop_str: str = None
    stop_token_ids: List[int] = None

    def get_prompt(self) -> str:
        """Get the prompt for generation."""
        system_prompt = self.system_template.format(system_message=self.system_message)
        if self.sep_style == SeparatorStyle.DeepSeek:
            seps = [self.sep, self.sep2]
            if system_prompt == "" or system_prompt is None:
                ret = ""
            else:
                ret = system_prompt + seps[0]
            for i, (role, message) in enumerate(self.messages):
                if message:
                    ret += role + ": " + message + seps[i % 2]
                else:
                    ret += role + ":"
            return ret
        elif self.sep_style == SeparatorStyle.DeepSeekV2:
            seps = [self.sep, self.sep2]
            if system_prompt == "" or system_prompt is None:
                ret = ""
            else:
                ret = system_prompt + seps[0]
            for i, (role, message) in enumerate(self.messages):
                if message:
                    if role == "User":
                        ret += "<｜sft▁begin｜>\n" + message + self.sep
                    else:
                        ret += message + self.sep2
                else:
                    ret = ret
            return ret
        elif self.sep_style == SeparatorStyle.PLAIN:
            seps = [self.sep, self.sep2]
            ret = ""
            for i, (role, message) in enumerate(self.messages):
                if message:
                    if type(message) is tuple:
                        message, _, _ = message
                    if i % 2 == 0:
                        ret += message + seps[i % 2]
                    else:
                        ret += message + seps[i % 2]
                else:
                    ret += ""
            return ret
        elif self.sep_style == SeparatorStyle.ALIGNMENT:
            seps = [self.sep, self.sep2]
            ret = ""
            for i, (role, message) in enumerate(self.messages):
                if message:
                    if type(message) is tuple:
                        message, _, _ = message
                    if i % 2 == 0:
                        ret += '<image>\n' + seps[i % 2]
                    else:
                        ret += message + seps[i % 2]
                else:
                    ret += ""
            return ret
        else:
            raise ValueError(f"Invalid style: {self.sep_style}")

    def set_system_message(self, system_message: str):
        """Set the system message."""
        self.system_message = system_message

    def append_message(self, role: str, message: str):
        """Append a new message."""
        self.messages.append([role, message])

    def update_last_message(self, message: str):
        """Update the last message."""
        self.messages[-1][1] = message

```

--------------------------------

### DeepseekV2DecoderLayer Initialization

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

Initializes the DeepseekV2DecoderLayer, configuring self-attention and the feed-forward network (MLP). It dynamically selects the attention implementation (eager, flash_attention_2, mla) based on the model's configuration and whether Multi-Layer Attention (MLA) is enabled. The MLP can be either a standard DeepseekV2MLP or a DeepseekV2MoE (Mixture of Experts) based on layer index and configuration.

```python
class DeepseekV2DecoderLayer(nn.Module):
    def __init__(self, config: DeepseekV2Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        
        if config.use_mla:
            attn_implementation = "mla_" + config._attn_implementation
        else:
            attn_implementation = "mha_" + config._attn_implementation
            
        self.self_attn = ATTENTION_CLASSES[attn_implementation](
            config=config, layer_idx=layer_idx
        )
        
        self.mlp = (
            DeepseekV2MoE(config)
            if (
                config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0
            )
            else DeepseekV2MLP(config)
        )
        self.input_layernorm = DeepseekV2RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_attention_layernorm = DeepseekV2RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
```

--------------------------------

### Configure Qwen2 Model for Deepseek-OCR-2

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet shows the configuration of a Qwen2 model with specific parameters for the Deepseek-OCR-2 project. It defines various dimensions, layer counts, attention heads, and other architectural details necessary for the model's operation. The configuration is then used to create a custom model instance.

```python
config = Qwen2Config(
  hidden_size=hidden_dimension,
  num_hidden_layers=decoder_layer,
  num_attention_heads=num_attention_heads,
  num_key_value_heads=num_key_value_heads,
  intermediate_size=intermediate_size,
  max_position_embeddings=max_position_embeddings,
  vocab_size=vocab_size,
  rms_norm_eps=rms_norm_eps,
  rope_theta=rope_theta,
  attention_dropout=attention_dropout,
  hidden_act=hidden_act,
  initializer_range=initializer_range,
  _attn_implementation=attn_implementation
)

self.model = self._create_custom_model(Qwen2Model, config)

del self.model.embed_tokens
```

--------------------------------

### Get Relative Positional Embeddings - PyTorch

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/18bdb3c4c8ccd6fe508f9fdedcd6438b01fcd3c5

This function calculates relative positional embeddings based on query and key sizes. It supports interpolation of embeddings if the provided relative position tensor's size does not match the required size, and scales coordinates based on query and key dimensions. It returns the extracted positional embeddings.

```python
def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
    """
    Get relative positional embeddings according to the relative positions of
    query and key sizes.
    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

    Returns:
        Extracted positional embeddings according to relative positions.
    """
    max_rel_dist = int(2 * max(q_size, k_size) - 1)
    # Interpolate rel pos if needed.
    if rel_pos.shape[0] != max_rel_dist:
        # Interpolate rel pos.
        dtype = rel_pos.dtype
        rel_pos = rel_pos.to(torch.float32)
        rel_pos_resized = F.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode="linear",
        ).to(dtype)
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)
    else:
        rel_pos_resized = rel_pos

    # Scale the coords with short length if shapes for q and k are different.
    q_coords = torch.arange(q_size, device=rel_pos.device)[:, None] * max(k_size / q_size, 1.0)
    k_coords = torch.arange(k_size, device=rel_pos.device)[None, :] * max(q_size / k_size, 1.0)
    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

    return rel_pos_resized[relative_coords.long()]
```

--------------------------------

### Window Partitioning Utility in Python

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/18bdb3c4c8ccd6fe508f9fdedcd6438b01fcd3c5

Provides functions for partitioning an input tensor into non-overlapping windows and then unpartitioning them back. This is useful for Swin Transformer-like architectures where computations are localized within windows. It handles padding to ensure windows are of a uniform size. Dependencies include PyTorch (torch, nn, F).

```python
def window_partition(x: torch.Tensor, window_size: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
    """
    Partition into non-overlapping windows with padding if needed.
    Args:
    x (tensor): input tokens with [B, H, W, C].
    window_size (int): window size.
    
    Returns:
    windows: windows after partition with [B * num_windows, window_size, window_size, C].
    (Hp, Wp): padded height and width before partition  
    """
    B, H, W, C = x.shape
    
    pad_h = (window_size - H % window_size) % window_size
    pad_w = (window_size - W % window_size) % window_size
    if pad_h > 0 or pad_w > 0:
        x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
    Hp, Wp = H + pad_h, W + pad_w
    
    x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows, (Hp, Wp)
    
def window_unpartition(
    windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
) -> torch.Tensor:
    """
    Window unpartition into original sequences and removing padding.
    Args:
    
```

--------------------------------

### Distributed Data Handling and Expert Processing in PyTorch

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/e6322a289fe5b5218278d276d4e7c58e8103f46a

This snippet details the process of distributing and gathering tokens across multiple expert ranks in a distributed deep learning setup. It involves `all_to_all` communication for token exchange and subsequent processing by individual experts. Dependencies include PyTorch (`torch`) and its distributed module (`dist`).

```python
input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()
dist.all_to_all(
    list(gathered_tokens.split(output_splits)),
    list(sorted_tokens.split(input_split_sizes)),
)
tokens_per_expert_post_gather = tokens_per_expert_group.view(
    self.ep_size, self.experts_per_rank
).sum(dim=0)
gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)
s = 0
for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):
    gatherd_idxs[s : s + k] = i % self.experts_per_rank
    s += k
gatherd_idxs = gatherd_idxs.argsort()
sorted_tokens = gathered_tokens[gatherd_idxs]
tokens_per_expert = tokens_per_expert_post_gather
tokens_per_expert = tokens_per_expert.cpu().numpy()

outputs = []
start_idx = 0
for i, num_tokens in enumerate(tokens_per_expert):
    end_idx = start_idx + num_tokens
    if num_tokens == 0:
        continue
    expert = self.experts[i + self.ep_rank * self.experts_per_rank]
    tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
    expert_out = expert(tokens_for_this_expert)
    outputs.append(expert_out)
    start_idx = end_idx

outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
if self.ep_size > 1:
    new_x = torch.empty_like(outs)
    new_x[gatherd_idxs] = outs
    gathered_tokens = new_x.new_empty(*sorted_tokens_shape)
    dist.all_to_all(
        list(gathered_tokens.split(input_split_sizes)),
        list(new_x.split(output_splits)),
    )
    outs = gathered_tokens

new_x = torch.empty_like(outs)
new_x[idxs] = outs
final_out = (
    new_x.view(*topk_ids.shape, -1)
    .type(topk_weight.dtype)
    .mul_(topk_weight.unsqueeze(dim=-1))
    .sum(dim=1)
    .type(new_x.dtype)
)
return final_out
```

--------------------------------

### Load DeepSeek-OCR-2 Model and Tokenizer with Transformers

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/raw/main/README

This snippet demonstrates how to load the DeepSeek-OCR-2 model and its corresponding tokenizer from Huggingface. It sets up the environment for GPU usage and specifies the model name. The model is configured to use flash attention for efficiency and loaded in bfloat16 precision.

```python
from transformers import AutoModel, AutoTokenizer
import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
model_name = 'deepseek-ai/DeepSeek-OCR-2'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, _attn_implementation='flash_attention_2', trust_remote_code=True, use_safetensors=True)
model = model.eval().cuda().to(torch.bfloat16)
```

--------------------------------

### Prepare OCR Prompts for DeepSeek-OCR 2 in Python

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/27738bf40dba31d9396930b2d01822303987ae47

This Python code defines various prompts for the DeepSeek-OCR 2 model to perform different OCR tasks. It includes prompts for converting documents to markdown, extracting text without layout information, and a basic free OCR prompt. These prompts are essential for guiding the model's output.

```python
# Prompt for markdown conversion with layout analysis
prompt_markdown = "<image>\n<|grounding|>Convert the document to markdown. "

# Prompt for free OCR without layout analysis
prompt_free_ocr = "<image>\nFree OCR. "
```

--------------------------------

### Initialize Deepseek-OCR Model Components (Python)

Source: https://huggingface.co/deepseek-ai/DeepSeek-OCR-2/commit/18bdb3c4c8ccd6fe508f9fdedcd6438b01fcd3c5

Initializes the core components of the Deepseek-OCR model, including patch embedding, positional embeddings, transformer blocks, and a convolutional neck. It configures attention mechanisms based on parameters like `global_attn_indexes` and `window_size`.

```python
self.patch_embed = PatchEmbed(
    kernel_size=(patch_size, patch_size),
    stride=(patch_size, patch_size),
    in_chans=in_chans,
    embed_dim=embed_dim,
)

self.pos_embed: Optional[nn.Parameter] = None
if use_abs_pos:
    # Initialize absolute positional embedding with pretrain image size.
    self.pos_embed = nn.Parameter(
        torch.zeros(1, img_size // patch_size, img_size // patch_size, embed_dim)
    )

self.blocks = nn.ModuleList()
for i in range(depth):
    block = Block(
        dim=embed_dim,
        num_heads=num_heads,
        mlp_ratio=mlp_ratio,
        qkv_bias=qkv_bias,
        norm_layer=norm_layer,
        act_layer=act_layer,
        use_rel_pos=use_rel_pos,
        rel_pos_zero_init=rel_pos_zero_init,
        window_size=window_size if i not in global_attn_indexes else 0,
        input_size=(img_size // patch_size, img_size // patch_size),
    )
    self.blocks.append(block)

self.neck = nn.Sequential(
    nn.Conv2d(
        embed_dim,
        out_chans,
        kernel_size=1,
        bias=False,
    ),
    LayerNorm2d(out_chans),
    nn.Conv2d(
        out_chans,
        out_chans,
        kernel_size=3,
        padding=1,
        bias=False,
    ),
    LayerNorm2d(out_chans),
)

self.net_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1, bias=False)
self.net_3 = nn.Conv2d(512, 896, kernel_size=3, stride=2, padding=1, bias=False)
```
Tokens:
Raw