tinydoc-vlm 0.2.0__py3-none-any.whl

tinydoc_vlm/modeling.py

@@ -0,0 +1,194 @@
+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", 49152)
+        
+        # 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)

tinydoc_vlm/output_heads.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+from typing import Dict
+
+class JSONHead(nn.Module):
+    """
+    Structured JSON generation head.
+    Projects decoder hidden states to schema-constrained JSON token logits.
+    In practice, this works as a specialized classifier over JSON structural tokens
+    plus a pointer network for field values.
+    """
+    def __init__(self, hidden_size: int, num_json_tokens: int = 256):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, num_json_tokens),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.proj(hidden_states)
+
+
+class KVHead(nn.Module):
+    """
+    Key-Value extraction head.
+    Produces key-value pairs from decoder hidden states with confidence scores.
+    Uses two separate projections: one for key detection, one for value extraction.
+    """
+    def __init__(self, hidden_size: int, num_keys: int = 128):
+        super().__init__()
+        self.key_classifier = nn.Linear(hidden_size, num_keys)
+        self.value_proj = nn.Linear(hidden_size, hidden_size)
+        self.confidence = nn.Linear(hidden_size, 1)
+
+    def forward(self, hidden_states: torch.Tensor) -> Dict[str, torch.Tensor]:
+        key_logits = self.key_classifier(hidden_states)
+        value_embeds = self.value_proj(hidden_states)
+        conf = torch.sigmoid(self.confidence(hidden_states))
+        return {
+            "key_logits": key_logits,
+            "value_embeds": value_embeds,
+            "confidence": conf,
+        }
+
+
+class TableHead(nn.Module):
+    """
+    Table generation head.
+    Outputs HTML/Markdown table tokens via a specialized projection.
+    Can also output structured row/cell coordinates.
+    """
+    def __init__(self, hidden_size: int, num_table_tokens: int = 128):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, num_table_tokens),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.proj(hidden_states)
+
+
+class OCRHead(nn.Module):
+    """
+    OCR text generation head.
+    Decodes visual features into character-level sequences.
+    Useful for reading-order text extraction.
+    """
+    def __init__(self, hidden_size: int, vocab_size: int = 256):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, vocab_size),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.proj(hidden_states)
+
+
+class QAHead(nn.Module):
+    """
+    Question-answering head.
+    Standard LM head for natural language answers.
+    """
+    def __init__(self, hidden_size: int, vocab_size: int):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, vocab_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.proj(hidden_states)
+
+
+class MultiTaskOutputHeads(nn.Module):
+    """
+    Container for all structured output heads.
+    Routes decoder hidden states to the appropriate task-specific head(s).
+    During training, all heads can be trained jointly with task-specific losses.
+    During inference, only the requested head is used.
+    """
+    def __init__(
+        self,
+        hidden_size: int,
+        vocab_size: int,
+        num_json_tokens: int = 256,
+        num_keys: int = 128,
+        num_table_tokens: int = 128,
+        ocr_vocab_size: int = 256,
+    ):
+        super().__init__()
+        self.json_head = JSONHead(hidden_size, num_json_tokens)
+        self.kv_head = KVHead(hidden_size, num_keys)
+        self.table_head = TableHead(hidden_size, num_table_tokens)
+        self.ocr_head = OCRHead(hidden_size, ocr_vocab_size)
+        self.qa_head = QAHead(hidden_size, vocab_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        task: str = "qa",
+    ) -> Dict[str, torch.Tensor]:
+        if task == "json":
+            return {"logits": self.json_head(hidden_states)}
+        elif task == "kv":
+            return self.kv_head(hidden_states)
+        elif task == "table":
+            return {"logits": self.table_head(hidden_states)}
+        elif task == "ocr":
+            return {"logits": self.ocr_head(hidden_states)}
+        elif task == "qa":
+            return {"logits": self.qa_head(hidden_states)}
+        else:
+            raise ValueError(f"Unknown task: {task}")

tinydoc_vlm/processing.py

@@ -0,0 +1,177 @@
+"""
+TinyDocVLMProcessor — standalone processor (does not inherit ProcessorMixin to avoid
+strict type-checking issues in transformers<4.45). Provides the same public API.
+"""
+import os
+import json
+from PIL import Image
+from typing import Dict, Any, Union, Optional, List
+
+import torch
+from transformers import AutoTokenizer
+
+from .image_processing import TinyDocImageProcessor
+
+
+class TinyDocVLMProcessor:
+    """
+    Coordinates TinyDocImageProcessor (image tiling + normalisation) and the
+    SmolLM2 tokenizer extended with document-special tokens.
+
+    Usage:
+        processor = TinyDocVLMProcessor()
+        inputs = processor(text=["Extract fields: <image>"], images=[pil_img])
+        # inputs → {"input_ids", "attention_mask", "pixel_values", "image_token_id"}
+    """
+
+    # Class-level attrs used by some HF utilities (save_pretrained etc.)
+    image_processor_class = "TinyDocImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor: Optional[TinyDocImageProcessor] = None,
+        tokenizer=None,
+        config=None,
+        **kwargs,
+    ):
+        self.image_processor = image_processor or TinyDocImageProcessor()
+        self.tokenizer = tokenizer or AutoTokenizer.from_pretrained(
+            "HuggingFaceTB/SmolLM2-135M-Instruct"
+        )
+        self.config = config
+
+        # Ensure <image> special token exists
+        self.image_token = "<image>"
+        if self.image_token not in self.tokenizer.get_vocab():
+            self.tokenizer.add_special_tokens(
+                {"additional_special_tokens": [self.image_token]}
+            )
+        self.image_token_id = self.tokenizer.convert_tokens_to_ids(self.image_token)
+
+    # ------------------------------------------------------------------
+    # Core __call__
+    # ------------------------------------------------------------------
+
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        images: Optional[Union[Image.Image, List[Image.Image]]] = None,
+        padding: bool = True,
+        truncation: bool = True,
+        max_length: Optional[int] = None,
+        return_tensors: str = "pt",
+    ) -> Dict[str, Any]:
+        """
+        Preprocesses text and images into tensors for the model.
+
+        Returns:
+            dict with keys: input_ids, attention_mask, pixel_values (optional),
+                            image_token_id.
+        """
+        # ---- 1. Image processing ------------------------------------------------
+        pixel_values = None
+        num_tiles_list: List[int] = []
+
+        if images is not None:
+            if not isinstance(images, list):
+                images = [images]
+
+            processed: List[torch.Tensor] = []
+            for img in images:
+                tile_tensor = self.image_processor.preprocess(img)  # (T, 3, H, W)
+                processed.append(tile_tensor)
+                num_tiles_list.append(tile_tensor.shape[0])
+
+            # Pad to max tiles so we can stack into a single tensor
+            max_tiles = max(num_tiles_list)
+            padded: List[torch.Tensor] = []
+            sz = self.image_processor.image_size
+            for tile_tensor in processed:
+                T = tile_tensor.shape[0]
+                if T < max_tiles:
+                    pad = torch.zeros(
+                        (max_tiles - T, 3, sz, sz),
+                        dtype=tile_tensor.dtype,
+                        device=tile_tensor.device,
+                    )
+                    tile_tensor = torch.cat([tile_tensor, pad], dim=0)
+                padded.append(tile_tensor)
+
+            pixel_values = torch.stack(padded, dim=0)  # (B, max_tiles, 3, H, W)
+
+        # ---- 2. Expand <image> tokens in text ----------------------------------
+        scale = (
+            getattr(self.config, "pixel_shuffle_scale", 3) if self.config else 3
+        )
+        patch_size = (
+            getattr(self.config, "patch_size", 16) if self.config else 16
+        )
+        sz = self.image_processor.image_size
+        tokens_per_tile = (sz // patch_size // scale) ** 2
+
+        if isinstance(text, list):
+            expanded: List[str] = []
+            for idx, t in enumerate(text):
+                if idx < len(num_tiles_list):
+                    total_vis = num_tiles_list[idx] * tokens_per_tile
+                else:
+                    total_vis = 0
+                expanded.append(self._expand_image_tokens(t, total_vis))
+            processed_text = expanded
+        else:
+            total_vis = num_tiles_list[0] * tokens_per_tile if num_tiles_list else 0
+            processed_text = self._expand_image_tokens(text, total_vis)
+
+        # ---- 3. Tokenise -------------------------------------------------------
+        enc = self.tokenizer(
+            processed_text,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+        )
+
+        inputs: Dict[str, Any] = {
+            "input_ids": enc["input_ids"],
+            "attention_mask": enc["attention_mask"],
+        }
+        if pixel_values is not None:
+            inputs["pixel_values"] = pixel_values
+        inputs["image_token_id"] = self.image_token_id
+
+        return inputs
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _expand_image_tokens(self, text: str, total_tokens: int) -> str:
+        """Replaces the single '<image>' placeholder with `total_tokens` copies."""
+        expansion = self.image_token * total_tokens
+        return text.replace(self.image_token, expansion)
+
+    # ------------------------------------------------------------------
+    # save / from_pretrained stubs (for HF Hub compatibility)
+    # ------------------------------------------------------------------
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        os.makedirs(save_directory, exist_ok=True)
+        self.image_processor.save_pretrained(save_directory)
+        self.tokenizer.save_pretrained(save_directory)
+        # Write a minimal processor_config.json
+        cfg = {
+            "processor_class": "TinyDocVLMProcessor",
+            "image_token": self.image_token,
+            "image_token_id": self.image_token_id,
+        }
+        with open(os.path.join(save_directory, "processor_config.json"), "w") as f:
+            json.dump(cfg, f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        image_processor = TinyDocImageProcessor.from_pretrained(
+            pretrained_model_name_or_path
+        )
+        tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
+        return cls(image_processor=image_processor, tokenizer=tokenizer, **kwargs)

tinydoc_vlm/token_compressor.py

@@ -0,0 +1,82 @@
+import math
+import torch
+import torch.nn as nn
+from .configuration import TinyDocVLMConfig
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        variance = x.pow(2).mean(-1, keepdim=True)
+        return x * torch.rsqrt(variance + self.eps) * self.weight
+
+class PixelShuffleTokenCompressor(nn.Module):
+    """
+    Performs space-to-depth token compression on Vision Transformer patch sequences.
+    Groups scale_factor x scale_factor patches and projects to decoder hidden dimension.
+    """
+    def __init__(self, config: TinyDocVLMConfig, encoder_dim: int, decoder_dim: int):
+        super().__init__()
+        self.config = config
+        self.scale_factor = config.pixel_shuffle_scale
+        self.encoder_dim = encoder_dim
+        self.decoder_dim = decoder_dim
+        
+        # After space-to-depth, channel dimension becomes encoder_dim * scale_factor^2
+        compressed_dim = encoder_dim * (self.scale_factor ** 2)
+        
+        # MLP projection to map visual tokens to language model dimension
+        self.projection = nn.Sequential(
+            nn.Linear(compressed_dim, decoder_dim),
+            nn.GELU(),
+            nn.Linear(decoder_dim, decoder_dim)
+        )
+        self.norm = RMSNorm(decoder_dim)
+
+    def forward(self, visual_features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            visual_features: shape (batch_size, num_tiles, num_patches, encoder_dim)
+            
+        Returns:
+            compressed_features: shape (batch_size, num_tiles, num_compressed_patches, decoder_dim)
+        """
+        batch_size, num_tiles, num_patches, encoder_dim = visual_features.shape
+        
+        # Determine spatial dimensions assuming a square grid of patches
+        grid_size = int(math.sqrt(num_patches))
+        if grid_size * grid_size != num_patches:
+            raise ValueError(
+                f"Number of patches ({num_patches}) must be a perfect square to apply 2D pixel shuffle."
+            )
+            
+        if grid_size % self.scale_factor != 0:
+            raise ValueError(
+                f"Grid size ({grid_size}) must be divisible by pixel_shuffle_scale ({self.scale_factor})."
+            )
+
+        # Reshape to 2D spatial grid: (batch_size * num_tiles, grid_size, grid_size, encoder_dim)
+        x = visual_features.view(batch_size * num_tiles, grid_size, grid_size, encoder_dim)
+        
+        # Apply space-to-depth: (batch_size * num_tiles, H//s, s, W//s, s, C)
+        s = self.scale_factor
+        x = x.view(batch_size * num_tiles, grid_size // s, s, grid_size // s, s, encoder_dim)
+        
+        # Permute: (batch_size * num_tiles, H//s, W//s, s, s, C)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        
+        # Reshape to flatten the spatial groups into the channel dimension:
+        # (batch_size * num_tiles, (H//s) * (W//s), s * s * C)
+        new_patches = (grid_size // s) ** 2
+        x = x.view(batch_size * num_tiles, new_patches, s * s * encoder_dim)
+        
+        # Project and normalize
+        x = self.projection(x)
+        x = self.norm(x)
+        
+        # Reshape back to batch: (batch_size, num_tiles, new_patches, decoder_dim)
+        x = x.view(batch_size, num_tiles, new_patches, self.decoder_dim)
+        return x