### 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: \n<|grounding|>Convert the document to markdown. # other image: \n<|grounding|>OCR this image. # without layouts: \nFree OCR. # figures in document: \nParse the figure. # general: \nDescribe this image in detail. # rec: \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: <|grounding|>Convert the document to markdown. # other image: <|grounding|>OCR this image. # without layouts: Free OCR. # figures in document: Parse the figure. # general: Describe this image in detail. # rec: 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 = " \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: \n<|grounding|>Convert the document to markdown. # without layouts: \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 += '\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 = "\n<|grounding|>Convert the document to markdown. " # Prompt for free OCR without layout analysis prompt_free_ocr = "\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) ```