|
|
from typing import List, Optional, Tuple, Union |
|
|
|
|
|
import torch |
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
from transformers.models.qwen2 import Qwen2ForCausalLM |
|
|
|
|
|
from .configuration_dots import DotsVisionConfig, DotsOCRConfig |
|
|
from .modeling_dots_vision import DotsVisionTransformer |
|
|
|
|
|
|
|
|
DOTS_VLM_MAX_IMAGES = 200 |
|
|
|
|
|
|
|
|
class DotsOCRForCausalLM(Qwen2ForCausalLM): |
|
|
config_class = DotsOCRConfig |
|
|
|
|
|
def __init__(self, config: DotsOCRConfig): |
|
|
super().__init__(config) |
|
|
|
|
|
if isinstance(self.config.vision_config, dict): |
|
|
vision_config = DotsVisionConfig(**self.config.vision_config) |
|
|
self.config.vision_config = vision_config |
|
|
else: |
|
|
vision_config = self.config.vision_config |
|
|
|
|
|
self.vision_tower = DotsVisionTransformer(vision_config) |
|
|
|
|
|
def prepare_inputs_embeds( |
|
|
self, |
|
|
input_ids: torch.LongTensor, |
|
|
pixel_values: Optional[torch.FloatTensor] = None, |
|
|
grid_thw: Optional[torch.FloatTensor] = None, |
|
|
img_mask: Optional[torch.BoolTensor] = None, |
|
|
) -> torch.Tensor: |
|
|
inputs_embeds = self.get_input_embeddings()(input_ids) |
|
|
|
|
|
if pixel_values is not None: |
|
|
assert img_mask is not None |
|
|
if grid_thw.shape[0] > DOTS_VLM_MAX_IMAGES: |
|
|
print( |
|
|
f"Num image exceeded: {grid_thw.shape[0]} > {DOTS_VLM_MAX_IMAGES}, which may cause FSDP hang" |
|
|
) |
|
|
|
|
|
vision_embeddings = self.vision_tower(pixel_values, grid_thw) |
|
|
|
|
|
true_indices = torch.nonzero(img_mask).squeeze() |
|
|
if len(true_indices) > vision_embeddings.size(0): |
|
|
print( |
|
|
f"img_mask sum > VE and will be truncated, mask.sum()={len(true_indices)} {vision_embeddings.size(0)=}" |
|
|
) |
|
|
true_indices = true_indices[: vision_embeddings.size(0)] |
|
|
new_img_mask = torch.zeros_like(img_mask, device=img_mask.device) |
|
|
new_img_mask[true_indices[:, 0], true_indices[:, 1]] = True |
|
|
else: |
|
|
new_img_mask = img_mask |
|
|
|
|
|
assert ( |
|
|
vision_embeddings.size(0) == new_img_mask.sum() |
|
|
), f"{vision_embeddings.size(0)=}, {new_img_mask.sum()=}" |
|
|
|
|
|
inputs_embeds = inputs_embeds.masked_scatter( |
|
|
new_img_mask.to(inputs_embeds.device).unsqueeze(-1).expand_as(inputs_embeds), |
|
|
vision_embeddings.to(inputs_embeds.device).type(inputs_embeds.dtype), |
|
|
) |
|
|
|
|
|
return inputs_embeds |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
input_ids: torch.LongTensor, |
|
|
pixel_values: Optional[torch.FloatTensor] = None, |
|
|
image_grid_thw: Optional[torch.FloatTensor] = None, |
|
|
inputs_embeds: Optional[torch.Tensor] = None, |
|
|
attention_mask: Optional[torch.Tensor] = None, |
|
|
position_ids: Optional[torch.LongTensor] = None, |
|
|
past_key_values: Optional[List[torch.FloatTensor]] = None, |
|
|
labels: Optional[torch.LongTensor] = None, |
|
|
output_attentions: Optional[bool] = None, |
|
|
output_hidden_states: Optional[bool] = None, |
|
|
return_dict: Optional[bool] = None, |
|
|
use_cache: Optional[bool] = None, |
|
|
logits_to_keep: int = 0, |
|
|
**loss_kwargs, |
|
|
) -> Union[Tuple, CausalLMOutputWithPast]: |
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
|
|
assert len(input_ids) >= 1, f"empty input_ids {input_ids.shape=} will cause gradnorm nan" |
|
|
if inputs_embeds is None: |
|
|
img_mask = input_ids == self.config.image_token_id |
|
|
inputs_embeds = self.prepare_inputs_embeds(input_ids, pixel_values, image_grid_thw, img_mask) |
|
|
|
|
|
outputs = super().forward( |
|
|
inputs_embeds=inputs_embeds, |
|
|
attention_mask=attention_mask, |
|
|
position_ids=position_ids, |
|
|
past_key_values=past_key_values, |
|
|
labels=labels, |
|
|
use_cache=use_cache if use_cache is not None else self.config.use_cache, |
|
|
output_attentions=output_attentions, |
|
|
output_hidden_states=output_hidden_states, |
|
|
|
|
|
logits_to_keep=logits_to_keep, |
|
|
**loss_kwargs, |
|
|
) |
|
|
|
|
|
return outputs |
|
|
|
|
|
def prepare_inputs_for_generation( |
|
|
self, |
|
|
input_ids, |
|
|
past_key_values=None, |
|
|
inputs_embeds=None, |
|
|
pixel_values=None, |
|
|
attention_mask=None, |
|
|
cache_position=None, |
|
|
num_logits_to_keep=None, |
|
|
**kwargs, |
|
|
): |
|
|
model_inputs = super().prepare_inputs_for_generation( |
|
|
input_ids, |
|
|
past_key_values=past_key_values, |
|
|
inputs_embeds=inputs_embeds, |
|
|
attention_mask=attention_mask, |
|
|
cache_position=cache_position, |
|
|
num_logits_to_keep=num_logits_to_keep, |
|
|
**kwargs, |
|
|
) |
|
|
|
|
|
if cache_position[0] == 0: |
|
|
model_inputs["pixel_values"] = pixel_values |
|
|
|
|
|
return model_inputs |
|
|
|
