tinydoc-vlm 0.2.0__py3-none-any.whl

tinydoc_vlm/__init__.py

@@ -0,0 +1,40 @@
+from .configuration import TinyDocVLMConfig
+from .vision_encoder import SigLIPVisionEncoder
+from .token_compressor import PixelShuffleTokenCompressor
+from .decoder import TinyDocDecoder
+from .modeling import TinyDocVLMForConditionalGeneration, TinyDocVLMPreTrainedModel
+from .image_processing import TinyDocImageProcessor
+from .processing import TinyDocVLMProcessor
+from .output_heads import MultiTaskOutputHeads, JSONHead, KVHead, TableHead, OCRHead, QAHead
+from .data import DocumentDataset
+from .losses import CombinedLoss
+from .trainer import TinyDocVLMTrainer, TrainerConfig
+
+__all__ = [
+    "TinyDocVLMConfig",
+    "SigLIPVisionEncoder",
+    "PixelShuffleTokenCompressor",
+    "TinyDocDecoder",
+    "TinyDocVLMForConditionalGeneration",
+    "TinyDocVLMPreTrainedModel",
+    "TinyDocImageProcessor",
+    "TinyDocVLMProcessor",
+    "MultiTaskOutputHeads",
+    "JSONHead",
+    "KVHead",
+    "TableHead",
+    "OCRHead",
+    "QAHead",
+    "DocumentDataset",
+    "CombinedLoss",
+    "TinyDocVLMTrainer",
+    "TrainerConfig",
+]
+
+from transformers import AutoConfig, AutoModelForCausalLM
+
+try:
+    AutoConfig.register("tinydoc_vlm", TinyDocVLMConfig)
+    AutoModelForCausalLM.register(TinyDocVLMConfig, TinyDocVLMForConditionalGeneration)
+except ValueError:
+    pass

tinydoc_vlm/attention.py

@@ -0,0 +1,76 @@
+import torch
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> torch.Tensor:
+    """
+    Generate 2D sinusoidal positional embeddings.
+    
+    Args:
+        embed_dim: Dimension of the embedding (must be even)
+        grid_size: Height/Width of the grid (assumed square)
+        cls_token: If True, prepends a zero embedding for the class token
+        
+    Returns:
+        pos_embed: shape (grid_size * grid_size, embed_dim) or (1 + grid_size * grid_size, embed_dim)
+    """
+    grid_h = torch.arange(grid_size, dtype=torch.float32)
+    grid_w = torch.arange(grid_size, dtype=torch.float32)
+    
+    # Create coordinate grid
+    grid = torch.meshgrid(grid_h, grid_w, indexing="ij")
+    grid = torch.stack(grid, dim=0)  # shape (2, grid_size, grid_size)
+    grid = grid.reshape(2, 1, grid_size, grid_size)
+    
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    
+    if cls_token:
+        # Prepend a zero embedding for CLS token
+        pos_embed = torch.cat([torch.zeros([1, embed_dim]), pos_embed], dim=0)
+        
+    return pos_embed
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: torch.Tensor) -> torch.Tensor:
+    """
+    Helper function to generate embeddings from a coordinate grid.
+    
+    Args:
+        embed_dim: Dimension of the embedding
+        grid: shape (2, 1, grid_h, grid_w) where index 0 is Y, index 1 is X
+        
+    Returns:
+        emb: shape (grid_h * grid_w, embed_dim)
+    """
+    assert embed_dim % 2 == 0
+    
+    # Use half of dimensions for X, half for Y
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # Y coords
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # X coords
+    
+    emb = torch.cat([emb_h, emb_w], dim=1)  # shape (grid_h * grid_w, embed_dim)
+    return emb
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: torch.Tensor) -> torch.Tensor:
+    """
+    Generate 1D sinusoidal positional embeddings for a coordinate tensor.
+    
+    Args:
+        embed_dim: Dimension of the embedding
+        pos: Coordinate tensor of shape (1, H, W) or (L,)
+        
+    Returns:
+        emb: shape (H * W, embed_dim) or (L, embed_dim)
+    """
+    assert embed_dim % 2 == 0
+    
+    # omega = 1 / (10000 ** (2i / d))
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / (10000 ** omega)  # shape (embed_dim // 2,)
+    
+    pos = pos.reshape(-1)  # Flatten spatial dims to 1D sequence
+    out = torch.outer(pos, omega)  # shape (seq_len, embed_dim // 2)
+    
+    emb_sine = torch.sin(out)
+    emb_cosine = torch.cos(out)
+    
+    emb = torch.cat([emb_sine, emb_cosine], dim=1)  # shape (seq_len, embed_dim)
+    return emb

tinydoc_vlm/configuration.py

@@ -0,0 +1,83 @@
+from typing import Dict, Any, Union
+from transformers import PretrainedConfig, AutoConfig
+
+class TinyDocVLMConfig(PretrainedConfig):
+    model_type = "tinydoc_vlm"
+    is_composition = True
+
+    def __init__(
+        self,
+        vision_config: Union[Dict[str, Any], PretrainedConfig] = None,
+        decoder_config: Union[Dict[str, Any], PretrainedConfig] = None,
+        pixel_shuffle_scale: int = 3,
+        image_size: int = 384,
+        patch_size: int = 16,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        
+        # Set defaults if not provided
+        if vision_config is None:
+            # Default SigLIP-B/16-like configuration (approx 93M parameters)
+            vision_config = {
+                "model_type": "siglip_vision_model",
+                "hidden_size": 768,
+                "intermediate_size": 3072,
+                "num_hidden_layers": 12,
+                "num_attention_heads": 12,
+                "patch_size": 16,
+                "image_size": 384,
+                "num_channels": 3,
+                "layer_norm_eps": 1e-6,
+            }
+            
+        if decoder_config is None:
+            # Default SmolLM2-135M-like configuration (approx 135M parameters)
+            decoder_config = {
+                "model_type": "llama",
+                "vocab_size": 49152,
+                "hidden_size": 576,
+                "intermediate_size": 1536,
+                "num_hidden_layers": 30,
+                "num_attention_heads": 9,
+                "num_key_value_heads": 3,
+                "max_position_embeddings": 8192,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 273000.0,
+                "attention_bias": False,
+            }
+
+        # Initialize config objects
+        if isinstance(vision_config, dict):
+            vision_config_copy = vision_config.copy()
+            vision_model_type = vision_config_copy.pop("model_type", "siglip_vision_model")
+            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config_copy)
+        else:
+            self.vision_config = vision_config
+
+        if isinstance(decoder_config, dict):
+            decoder_config_copy = decoder_config.copy()
+            decoder_model_type = decoder_config_copy.pop("model_type", "llama")
+            self.decoder_config = AutoConfig.for_model(decoder_model_type, **decoder_config_copy)
+        else:
+            self.decoder_config = decoder_config
+
+        self.pixel_shuffle_scale = pixel_shuffle_scale
+        self.image_size = image_size
+        self.patch_size = patch_size
+
+    def __getattr__(self, name):
+        if name in ('decoder_config', 'vision_config'):
+            raise AttributeError(name)
+        if 'decoder_config' in self.__dict__:
+            try:
+                return getattr(self.decoder_config, name)
+            except AttributeError:
+                pass
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def to_dict(self) -> Dict[str, Any]:
+        output = super().to_dict()
+        output["vision_config"] = self.vision_config.to_dict()
+        output["decoder_config"] = self.decoder_config.to_dict()
+        return output

tinydoc_vlm/data.py

@@ -0,0 +1,92 @@
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+
+from .image_processing import TinyDocImageProcessor
+
+
+class DocumentDataset(Dataset):
+    """
+    Dataset for document understanding training.
+    Supports loading from a JSON manifest file or from individual samples.
+
+    Manifest format (JSONL):
+        {"image_path": "path/to/image.png", "text": "Extract: <image>", "labels": {...}}
+    """
+    def __init__(
+        self,
+        data_root: Union[str, Path],
+        manifest_path: Optional[Union[str, Path]] = None,
+        image_processor: Optional[TinyDocImageProcessor] = None,
+        max_seq_length: int = 2048,
+        stage: int = 1,
+        samples: Optional[List[Dict]] = None,
+    ):
+        self.data_root = Path(data_root)
+        self.image_processor = image_processor or TinyDocImageProcessor()
+        self.max_seq_length = max_seq_length
+        self.stage = stage
+
+        if samples is not None:
+            self.samples = samples
+        elif manifest_path:
+            with open(manifest_path) as f:
+                self.samples = [json.loads(line) for line in f if line.strip()]
+        else:
+            self.samples = []
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> Dict:
+        sample = self.samples[idx]
+        image_path = self.data_root / sample["image_path"]
+        image = Image.open(image_path).convert("RGB")
+        pixel_values = self.image_processor.preprocess(image)
+
+        text = sample.get("text", "<image>")
+        labels = sample.get("labels", {})
+
+        return {
+            "pixel_values": pixel_values,
+            "text": text,
+            "labels": labels,
+            "metadata": sample.get("metadata", {}),
+        }
+
+
+def collate_fn(batch: List[Dict], tokenizer, image_token_id: int, max_length: int = 2048) -> Dict:
+    """
+    Collate function for DocumentDataset.
+    Handles variable-length text, variable-number tiles, and label padding.
+    """
+    texts = [item["text"] for item in batch]
+    images = [item.get("pixel_values") for item in batch]
+
+    max_tiles = max(pv.shape[0] for pv in images)
+    image_size = images[0].shape[-1]
+    padded_pixel_values = []
+    for pv in images:
+        num_tiles = pv.shape[0]
+        if num_tiles < max_tiles:
+            pad = torch.zeros(max_tiles - num_tiles, 3, image_size, image_size, dtype=pv.dtype)
+            pv = torch.cat([pv, pad], dim=0)
+        padded_pixel_values.append(pv)
+
+    pixel_values = torch.stack(padded_pixel_values, dim=0)
+
+    tokenized = tokenizer(texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt")
+
+    labels = tokenized["input_ids"].clone()
+    labels[labels == tokenizer.pad_token_id] = -100
+
+    return {
+        "input_ids": tokenized["input_ids"],
+        "attention_mask": tokenized["attention_mask"],
+        "pixel_values": pixel_values,
+        "labels": labels,
+    }

tinydoc_vlm/decoder.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+from typing import Optional, List
+from transformers import LlamaForCausalLM, LlamaConfig
+
+class TinyDocDecoder(nn.Module):
+    """
+    Decoder wrapper around LlamaForCausalLM (used by SmolLM2).
+    Manages loading and vocabulary/embedding resizing for special tokens.
+    """
+    def __init__(self, config: LlamaConfig):
+        super().__init__()
+        self.config = config
+        self.lm = LlamaForCausalLM(config)
+        self.hidden_size = config.hidden_size
+
+    def resize_token_embeddings(self, new_num_tokens: int) -> nn.Embedding:
+        """
+        Resizes input token embeddings and output LM head of the decoder.
+        """
+        resized = self.lm.resize_token_embeddings(new_num_tokens)
+        self.config.vocab_size = new_num_tokens
+        return resized
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.lm.get_input_embeddings()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = 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,
+    ):
+        return self.lm(
+            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,
+        )

tinydoc_vlm/image_processing.py

@@ -0,0 +1,117 @@
+import numpy as np
+from PIL import Image
+from typing import Optional, List
+import torch
+import torchvision.transforms as T
+from transformers.image_processing_base import ImageProcessingMixin
+
+class TinyDocImageProcessor(ImageProcessingMixin):
+    """
+    Image processor for TinyDoc-VLM.
+    Handles resizing, normalization, and optional tiling (splitting) of document images.
+    """
+    def __init__(
+        self,
+        image_size: int = 384,
+        mean: Optional[List[float]] = None,
+        std: Optional[List[float]] = None,
+        tiling_mode: str = "auto",  # "none", "auto" (split if large)
+        **kwargs,
+    ):
+        self.image_size = image_size
+        self.mean = mean or [0.5, 0.5, 0.5]
+        self.std = std or [0.5, 0.5, 0.5]
+        self.tiling_mode = tiling_mode
+
+        super().__init__(**kwargs)
+        
+        # Base torchvision transforms for single tile
+        self.transform = T.Compose([
+            T.ToTensor(),
+            T.Normalize(mean=self.mean, std=self.std)
+        ])
+
+    def preprocess(
+        self, 
+        image: Image.Image, 
+        return_tensors: str = "pt"
+    ) -> torch.Tensor:
+        """
+        Preprocesses a PIL Image into a multi-tile float tensor.
+        
+        Returns shape: (num_tiles, 3, image_size, image_size)
+        """
+        # Ensure RGB
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+            
+        w, h = image.size
+        
+        if self.tiling_mode == "none" or (w <= self.image_size and h <= self.image_size):
+            # No tiling needed: resize to image_size x image_size and return single tile
+            resized = image.resize((self.image_size, self.image_size), Image.Resampling.BILINEAR)
+            tile_tensor = self.transform(resized)  # shape (3, image_size, image_size)
+            # Add tile batch dimension: shape (1, 3, image_size, image_size)
+            return tile_tensor.unsqueeze(0)
+            
+        # Tiling mode 'auto': split high-res image into a grid of image_size x image_size tiles,
+        # plus a downscaled overview thumbnail.
+        
+        # Calculate how many tiles we need
+        cols = int(np.ceil(w / self.image_size))
+        rows = int(np.ceil(h / self.image_size))
+        
+        # Limit grid size to prevent excessive memory usage (max 2x2 grid = 4 tiles)
+        cols = min(cols, 2)
+        rows = min(rows, 2)
+        
+        # Target size for the tiling grid
+        target_w = cols * self.image_size
+        target_h = rows * self.image_size
+        
+        # Resize original image to fit the target grid shape (maintaining proportions via padding)
+        resized_full = self._resize_and_pad(image, target_w, target_h)
+        
+        tiles = []
+        
+        # 1. Generate thumbnail/overview of the full image
+        thumbnail = image.resize((self.image_size, self.image_size), Image.Resampling.BILINEAR)
+        tiles.append(self.transform(thumbnail))
+        
+        # 2. Extract tiles from the grid
+        for r in range(rows):
+            for c in range(cols):
+                box = (
+                    c * self.image_size,
+                    r * self.image_size,
+                    (c + 1) * self.image_size,
+                    (r + 1) * self.image_size
+                )
+                tile = resized_full.crop(box)
+                tiles.append(self.transform(tile))
+                
+        # Stack tiles along a new dimension: shape (num_tiles, 3, image_size, image_size)
+        # where num_tiles = 1 (overview) + rows * cols
+        stacked_tiles = torch.stack(tiles, dim=0)
+        return stacked_tiles
+
+    def _resize_and_pad(self, img: Image.Image, target_w: int, target_h: int) -> Image.Image:
+        """
+        Resizes and pads an image to target dimensions while maintaining aspect ratio.
+        """
+        # Calculate aspect ratio
+        w, h = img.size
+        ratio = min(target_w / w, target_h / h)
+        new_w = int(w * ratio)
+        new_h = int(h * ratio)
+        
+        resized = img.resize((new_w, new_h), Image.Resampling.BILINEAR)
+        
+        # Create a new padded background image
+        padded = Image.new("RGB", (target_w, target_h), (255, 255, 255))
+        # Center the resized image
+        x_offset = (target_w - new_w) // 2
+        y_offset = (target_h - new_h) // 2
+        padded.paste(resized, (x_offset, y_offset))
+        
+        return padded

tinydoc_vlm/losses.py

@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional
+
+
+def ce_loss(logits: torch.Tensor, labels: torch.Tensor, ignore_index: int = -100) -> torch.Tensor:
+    return F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=ignore_index)
+
+
+def kv_loss(kv_outputs: Dict[str, torch.Tensor], key_labels: torch.Tensor, confidence_labels: torch.Tensor) -> torch.Tensor:
+    key_ce = F.cross_entropy(kv_outputs["key_logits"].view(-1, kv_outputs["key_logits"].size(-1)), key_labels.view(-1))
+    conf_bce = F.binary_cross_entropy(kv_outputs["confidence"].view(-1), confidence_labels.view(-1))
+    return key_ce + conf_bce
+
+
+class CombinedLoss(nn.Module):
+    """
+    Combines multiple task-specific losses for multi-stage training.
+    Stage 1: layout + OCR + region
+    Stage 2: QA + JSON + KV + table
+    Stage 3: standard LM
+    """
+    def __init__(self, stage: int = 1):
+        super().__init__()
+        self.stage = stage
+
+    def forward(
+        self,
+        lm_logits: torch.Tensor,
+        lm_labels: torch.Tensor,
+        head_outputs: Optional[Dict[str, torch.Tensor]] = None,
+        head_labels: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        losses = {}
+        total_loss = torch.tensor(0.0, device=lm_logits.device)
+
+        lm_loss = ce_loss(lm_logits, lm_labels)
+        losses["lm_loss"] = lm_loss
+
+        if self.stage == 1:
+            total_loss = lm_loss
+        elif self.stage == 2:
+            total_loss = lm_loss
+            if head_outputs and head_labels:
+                if "json_logits" in head_outputs and "json_labels" in head_labels:
+                    json_loss = ce_loss(head_outputs["json_logits"], head_labels["json_labels"])
+                    losses["json_loss"] = json_loss
+                    total_loss = total_loss + json_loss
+
+                if "kv" in head_outputs and "kv_key_labels" in head_labels:
+                    kv = kv_loss(head_outputs["kv"], head_labels["kv_key_labels"], head_labels["kv_conf_labels"])
+                    losses["kv_loss"] = kv
+                    total_loss = total_loss + kv
+
+                if "table_logits" in head_outputs and "table_labels" in head_labels:
+                    table_loss = ce_loss(head_outputs["table_logits"], head_labels["table_labels"])
+                    losses["table_loss"] = table_loss
+                    total_loss = total_loss + table_loss
+        else:
+            total_loss = lm_loss
+
+        losses["loss"] = total_loss
+        return losses