| from copy import deepcopy |
| from functools import partial |
| from unittest.mock import patch |
|
|
| import torch |
| from torch import nn |
| from sglang.srt.models import deepseek_ocr as vision |
| from sglang.srt.models import unlimited_ocr as base |
| from sglang.srt.models import unlimited_ocr_v2 as dash |
| from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear |
| from sglang.srt.models.unlimited_ocr_v2 import UnlimitedOCRV2ForCausalLM as UpstreamModel |
|
|
|
|
| class ConfigAttention(dash.UnlimitedOCRV2Attention): |
| def __init__(self, config, layer_id, quant_config=None, prefix=""): |
| |
| attention_config = deepcopy(config) |
| attention_config.hidden_size = config.num_attention_heads * config.head_dim |
| super().__init__(attention_config, layer_id, quant_config, prefix) |
| self.hidden_size = config.hidden_size |
| outputs = { |
| "q_proj": config.num_attention_heads * self.head_dim, |
| "k_proj": config.num_key_value_heads * self.head_dim, |
| "v_proj": config.num_key_value_heads * self.head_dim, |
| "index_q_proj": self.idx_head_dim, |
| "index_k_proj": self.idx_head_dim, |
| } |
| for name, output_size in outputs.items(): |
| setattr(self, name, ColumnParallelLinear(config.hidden_size, output_size, |
| bias=False, gather_output=name.startswith("index_"), quant_config=quant_config, |
| prefix=f"{prefix}.{name}" if prefix else name)) |
| self.o_proj = RowParallelLinear(config.num_attention_heads * self.head_dim, |
| config.hidden_size, bias=False, quant_config=quant_config, |
| prefix=f"{prefix}.o_proj" if prefix else "o_proj") |
|
|
|
|
| class UnlimitedOCRV2ForCausalLM(UpstreamModel): |
| def __init__(self, *, config, quant_config=None, prefix=""): |
| specs = config.vision_config.width |
|
|
| def sam(): |
| spec = specs["sam_vit_b"] |
| encoder = vision.ImageEncoderViT(embed_dim=spec["width"], depth=spec["layers"], |
| num_heads=spec["heads"], global_attn_indexes=spec["global_attn_indexes"], |
| out_chans=spec["neck_channels"], window_size=14, use_rel_pos=True, |
| norm_layer=partial(nn.LayerNorm, eps=1e-6)) |
| c1, c2 = spec["downsample_channels"] |
| encoder.net_2 = nn.Conv2d(spec["neck_channels"], c1, 3, stride=2, padding=1, bias=False) |
| encoder.net_3 = nn.Conv2d(c1, c2, 3, stride=2, padding=1, bias=False) |
| return encoder |
|
|
| def clip(): |
| spec = specs["clip-l-14-224"] |
| cfg = deepcopy(vision.vit_model_cfg) |
| cfg.update(num_layers=spec["layers"], hidden_size=spec["width"], num_heads=spec["heads"], |
| num_attention_heads=spec["heads"], ffn_hidden_size=spec["width"] * 4) |
| return vision.VitModel(cfg) |
|
|
| with (patch.object(base, "build_sam_vit_b", sam), |
| patch.object(base, "build_clip_l", clip), |
| patch.object(dash, "UnlimitedOCRV2Attention", ConfigAttention)): |
| super().__init__(config=config, quant_config=quant_config, prefix=prefix) |
|
|
|
|
| EntryClass = UnlimitedOCRV2ForCausalLM |
|
|
