File size: 5,011 Bytes
bb5b74a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from typing import Any, List, Optional
from transformers.configuration_utils import PretrainedConfig
from transformers import Qwen2_5_VLProcessor, AutoProcessor
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
class MonkeyOCRv2ViTAEEncoderConfig(PretrainedConfig):
model_type: str = "monkeyocrv2_vitae_encoder"
def __init__(
self,
num_channels: int = 3,
patch_size: int = 32, # preprocessor tile size in pixels (= encoder stride for 1:1 mapping)
temporal_patch_size: int = 1,
# ViTAEv2-S parameters (see Table 1 in paper):
stage_dims: List[int] = None, # [64, 128, 256, 512] token_dims
stage_depths: List[int] = None, # [2, 2, 8, 2] NC_depth (official ViTAEv2-S)
stage_heads: List[int] = None, # [1, 2, 4, 8] NC_heads
downsample_ratios: List[int] = None, # [4, 2, 2, 2] RC downsample ratio per stage
kernel_sizes: List[int] = None, # [7, 3, 3, 3] RC kernel size per stage
rc_tokens_type: List[str] = None, # ['window', 'window', 'transformer', 'transformer']
nc_tokens_type: List[str] = None, # ['window', 'window', 'transformer', 'transformer']
nc_groups: List[int] = None, # [1, 32, 64, 128] NC_group
rc_groups: List[int] = None, # [1, 16, 32, 64] RC PCM group (RC0→RC3)
rc_heads: List[int] = None, # [1, 1, 2, 4] RC attn heads (RC0→RC3)
rc_embed_dims: List[int] = None, # [64, 64, 128, 256] PRM/PCM intermediate dim per RC
prm_embed_dim: int = 64, # legacy fallback for RC0 PRM dim
window_size: int = 7, # windowed-attn window size (stages 1-2 + RC1-RC2)
mlp_ratio: float = 4.0,
hidden_size: int = 1024, # LLM projection output dim
rms_norm_eps: float = 1e-5,
use_bias: bool = False,
attn_implementation: str = "sdpa", # "eager"/"sdpa"/"flash_attention_2"
initializer_range: float = 0.02,
init_merger_std: float = 0.02,
is_causal: bool = False,
post_norm: bool = True,
gradient_checkpointing: bool = False,
**kwargs: Any,
):
super().__init__(**kwargs)
self.num_channels = num_channels
self.patch_size = patch_size
self.temporal_patch_size = temporal_patch_size
self.stage_dims = stage_dims if stage_dims is not None else [64, 128, 256, 512]
self.stage_depths = stage_depths if stage_depths is not None else [2, 2, 8, 2]
self.stage_heads = stage_heads if stage_heads is not None else [1, 2, 4, 8]
self.downsample_ratios = downsample_ratios if downsample_ratios is not None else [4, 2, 2, 2]
self.kernel_sizes = kernel_sizes if kernel_sizes is not None else [7, 3, 3, 3]
self.rc_tokens_type = rc_tokens_type if rc_tokens_type is not None else ["window", "window", "transformer", "transformer"]
self.nc_tokens_type = nc_tokens_type if nc_tokens_type is not None else ["window", "window", "transformer", "transformer"]
self.nc_groups = nc_groups if nc_groups is not None else [1, 32, 64, 128]
self.rc_groups = rc_groups if rc_groups is not None else [1, 16, 32, 64]
self.rc_heads = rc_heads if rc_heads is not None else [1, 1, 2, 4]
self.rc_embed_dims = rc_embed_dims if rc_embed_dims is not None else [64, 64, 128, 256]
self.prm_embed_dim = prm_embed_dim
self.window_size = window_size
self.mlp_ratio = mlp_ratio
self.hidden_size = hidden_size
self.rms_norm_eps = rms_norm_eps
self.use_bias = use_bias
self.attn_implementation = attn_implementation
self.initializer_range = initializer_range
self.init_merger_std = init_merger_std
self.is_causal = is_causal
self.post_norm = post_norm
self.gradient_checkpointing = gradient_checkpointing
class MonkeyOCRv2ViTAEProcessor(Qwen2_5_VLProcessor):
attributes = ["image_processor", "tokenizer"]
def __init__(self, image_processor=None, tokenizer=None, chat_template=None,
encoder_stride=None, **kwargs):
super().__init__(image_processor, tokenizer, chat_template=chat_template)
self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
self.image_token_id = 151655 if not hasattr(tokenizer, "image_token_id") else tokenizer.image_token_id
# encoder_stride: total conv downsampling applied to each preprocessor tile.
# With encoder_stride == preprocessor.patch_size, stage4 outputs 1 token per tile.
pp_size = getattr(image_processor, 'patch_size', 32) if image_processor is not None else 32
self.encoder_stride = encoder_stride if encoder_stride is not None else pp_size
AutoProcessor.register("monkeyocrv2_vitae", MonkeyOCRv2ViTAEProcessor)
CONFIG_MAPPING.register("monkeyocrv2_vitae_encoder", MonkeyOCRv2ViTAEEncoderConfig)
|