import torch import torch.nn as nn from typing import Dict, Any, Optional, List, Tuple, Union from transformers import PreTrainedModel, GenerationMixin from transformers.modeling_outputs import CausalLMOutputWithPast from .configuration import TinyDocVLMConfig from .vision_encoder import SigLIPVisionEncoder from .token_compressor import PixelShuffleTokenCompressor from .decoder import TinyDocDecoder from .output_heads import MultiTaskOutputHeads class TinyDocVLMPreTrainedModel(PreTrainedModel): config_class = TinyDocVLMConfig base_model_prefix = "tinydoc_vlm" supports_gradient_checkpointing = True def _init_weights(self, module): std = getattr(self.config, "initializer_range", 0.02) if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class TinyDocVLMForConditionalGeneration(TinyDocVLMPreTrainedModel, GenerationMixin): """ TinyDoc-VLM: The World's Smallest Document Understanding Model. Coordinates SigLIP Vision Encoder, PixelShuffle Compressor, and SmolLM2 Decoder. """ def __init__(self, config: TinyDocVLMConfig): super().__init__(config) # 1. Vision Encoder self.vision_encoder = SigLIPVisionEncoder(config) # 2. Token Compressor / Connector self.compressor = PixelShuffleTokenCompressor( config, encoder_dim=config.vision_config.hidden_size, decoder_dim=config.decoder_config.hidden_size ) # 3. Decoder self.decoder = TinyDocDecoder(config.decoder_config) # Learnable image pad / placeholder token ID self.image_token_id = getattr(config, "image_token_id", 49153) # 2D Positional Embeddings for visual features (added to tokens before projection) s = config.pixel_shuffle_scale compressed_grid_size = (config.image_size // config.patch_size) // s compressed_patches = compressed_grid_size ** 2 # Learnable 2D positional embeddings for the compressed visual tokens self.visual_pos_embed = nn.Parameter( torch.zeros(1, 1, compressed_patches, config.decoder_config.hidden_size) ) # 4. Structured Output Heads (multi-task) self.output_heads = MultiTaskOutputHeads( hidden_size=config.decoder_config.hidden_size, vocab_size=config.decoder_config.vocab_size, ) # Initialize weights self.post_init() def get_input_embeddings(self) -> nn.Module: return self.decoder.get_input_embeddings() def set_input_embeddings(self, value): self.decoder.lm.set_input_embeddings(value) def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, task: Optional[str] = None, ) -> Union[Tuple, Dict, CausalLMOutputWithPast]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict output_hidden_states = True if task else output_hidden_states # Decoding pass (no new visual input, reuse cached states) if pixel_values is None and past_key_values is not None: outputs = self.decoder( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) if task: hidden = outputs.hidden_states[-1] if hasattr(outputs, "hidden_states") else outputs[2] head_outputs = self.output_heads(hidden, task=task) return {"lm_outputs": outputs, "head_outputs": head_outputs} return outputs # Prefill pass: merge text and visual tokens into inputs_embeds if inputs_embeds is None: inputs_embeds = self.decoder.get_input_embeddings()(input_ids) if pixel_values is not None: visual_features = self.vision_encoder(pixel_values) compressed_features = self.compressor(visual_features) compressed_features = compressed_features + self.visual_pos_embed batch_size, num_tiles, compressed_patches, decoder_dim = compressed_features.shape flat_visual_features = compressed_features.view( batch_size, num_tiles * compressed_patches, decoder_dim ) image_mask = (input_ids == self.image_token_id) for b in range(batch_size): num_places = image_mask[b].sum().item() if num_places > 0: features_to_insert = flat_visual_features[b][:num_places] inputs_embeds[b, image_mask[b]] = features_to_insert outputs = self.decoder( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) if task: hidden = outputs.hidden_states[-1] if hasattr(outputs, "hidden_states") else outputs[-1] head_outputs = self.output_heads(hidden, task=task) return {"lm_outputs": outputs, "head_outputs": head_outputs} return outputs def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, pixel_values=None, **kwargs ) -> Dict[str, Any]: """ Overridden to support KV caching during auto-regressive generation. """ is_decoding = past_key_values is not None and pixel_values is None if is_decoding: input_ids = input_ids[:, -1:] inputs_embeds = None position_ids = kwargs.get("position_ids", None) if attention_mask is not None and position_ids is None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) if is_decoding: position_ids = position_ids[:, -input_ids.shape[-1]:] return { "input_ids": input_ids, "inputs_embeds": inputs_embeds, "past_key_values": past_key_values, "pixel_values": pixel_values, "attention_mask": attention_mask, "position_ids": position_ids, "use_cache": kwargs.get("use_cache"), } def _reorder_cache(self, past_key_values, beam_idx): return self.decoder.lm._reorder_cache(past_key_values, beam_idx)