Add text encoder config

Author: VikParuchuri

surya/layout.py

@@ -91,6 +91,8 @@ def batch_layout_detection(images: List, model, processor, batch_size=None) -> L
 
         decoder_position_ids = torch.ones_like(batch_decoder_input[0, :, 0], dtype=torch.int64, device=model.device).cumsum(0) - 1
         model.decoder.model._setup_cache(model.config, batch_size, model.device, model.dtype)
+        if hasattr(model, "text_encoder"):
+            model.text_encoder.model._setup_cache(model.config, batch_size, model.device, model.dtype)
 
         batch_predictions = [[] for _ in range(len(images))]
 
@@ -100,6 +102,24 @@ def batch_layout_detection(images: List, model, processor, batch_size=None) -> L
             token_count = 0
             all_done = torch.zeros(current_batch_size, dtype=torch.bool, device=model.device)
 
+            if hasattr(model, "text_encoder"):
+                text_encoder_input_ids = torch.arange(
+                    model.text_encoder.config.query_token_count,
+                    device=encoder_hidden_states.device,
+                    dtype=torch.long
+                ).unsqueeze(0).expand(encoder_hidden_states.size(0), -1)
+
+                text_encoder_hidden_states = model.text_encoder(
+                    input_ids=text_encoder_input_ids,
+                    cache_position=None,
+                    attention_mask=None,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=None,
+                    use_cache=False
+                ).hidden_states
+
+                encoder_hidden_states = torch.cat([encoder_hidden_states, text_encoder_hidden_states], dim=1)
+
             while token_count < settings.LAYOUT_MAX_BOXES:
                 is_prefill = token_count == 0
                 return_dict = model.decoder(

surya/model/layout/config.py

@@ -7,7 +7,7 @@
 from surya.settings import settings
 
 SPECIAL_TOKENS = 3
-QUERY_TOKENS = 192
+QUERY_TOKENS = 144
 BBOX_SIZE = 1024
 PADDED_BBOX_SIZE = BBOX_SIZE + 1
 
@@ -45,9 +45,14 @@ def __init__(self, **kwargs):
 
         encoder_config = kwargs.pop("encoder")
         decoder_config = kwargs.pop("decoder")
+        text_encoder_config = kwargs.pop("text_encoder", None)
 
         self.encoder = encoder_config
         self.decoder = decoder_config
+
+        if text_encoder_config is not None:
+            self.text_encoder = text_encoder_config
+
         self.is_encoder_decoder = True
 
         if isinstance(decoder_config, dict):
@@ -221,6 +226,91 @@ def __init__(
             **kwargs,
         )
 
+    @property
+    def layers_block_type(self):
+        return (self.block_types * 100)[: self.num_hidden_layers]
+
+
+class SuryaLayoutTextEncoderConfig(PretrainedConfig):
+    model_type = "surya_layout"
+
+    def __init__(
+        self,
+        num_hidden_layers=4,
+        vocab_size=256,
+        hidden_size=512,
+        intermediate_size=4 * 512,
+        encoder_hidden_size=1024,
+        num_attention_heads=8,
+        lru_width=None,
+        attention_window_size=16,
+        conv1d_width=4,
+        logits_soft_cap=30.0,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        bos_token_id=1,
+        hidden_activation="gelu_pytorch_tanh",
+        rope_theta=10000.0,
+        block_types=("attention",),
+        cross_attn_layers=(0, 1, 2, 3),
+        self_attn_layers=(0, 1, 2, 3),
+        global_attn_layers=(0, 1, 2, 3),
+        attention_dropout=0.0,
+        num_key_value_heads=4,
+        attention_bias=False,
+        w_init_variance_scale=0.01,
+        init_std=0.02,
+        tie_word_embeddings=False,
+        aux_heads=0, # How many n-token-ahead heads to add
+        iteration_count=1,
+        causal=False,
+        query_token_count=QUERY_TOKENS,
+        layer_norm_eps=1e-5,
+        **kwargs,
+    ):
+        self.num_hidden_layers = num_hidden_layers
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.lru_width = lru_width if lru_width is not None else hidden_size
+        self.attention_window_size = attention_window_size
+        self.conv1d_width = conv1d_width
+        self.logits_soft_cap = logits_soft_cap
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.block_types = list(block_types)
+        self.hidden_activation = hidden_activation
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+        if self.num_key_value_heads > self.num_attention_heads:
+            raise ValueError("The number of `num_key_value_heads` must be smaller than `num_attention_heads`")
+        self.cross_attn_layers = cross_attn_layers
+        self.self_attn_layers = self_attn_layers
+        self.global_attn_layers = global_attn_layers
+        self.attention_dropout = attention_dropout
+        self.attention_bias = attention_bias
+        self.w_init_variance_scale = w_init_variance_scale
+        self.final_w_init_variance_scale = 2.0 / self.num_hidden_layers
+        self.init_std = init_std
+        self.tie_word_embeddings = tie_word_embeddings
+        self.aux_heads = aux_heads
+        self.encoder_hidden_size = encoder_hidden_size
+        self.iteration_count = iteration_count
+        self.causal = causal
+        self.query_token_count = query_token_count
+        self.layer_norm_eps = layer_norm_eps
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+
     @property
     def layers_block_type(self):
         return (self.block_types * 100)[: self.num_hidden_layers]

surya/model/layout/decoder.py

@@ -8,7 +8,7 @@
 from torch import nn
 from torch.nn import functional as F
 
-from surya.model.common.adetr.decoder import SuryaADETRDecoderModel, SuryaADETRDecoderPreTrainedModel
+from surya.model.common.adetr.decoder import SuryaADETRDecoderModel, SuryaADETRDecoderPreTrainedModel, WrappedEmbedding
 from surya.model.layout.config import LayoutModelOutput
 from transformers.modeling_outputs import CausalLMOutput
 from surya.settings import settings
@@ -126,4 +126,63 @@ def forward(
             bbox_logits=bbox_logits,
             class_logits=class_logits,
             hidden_states=outputs.hidden_states,
+        )
+
+@dataclass
+class TextEncoderOutput(CausalLMOutput):
+    hidden_states: torch.FloatTensor = None
+
+
+class SuryaLayoutTextEncoder(SuryaADETRDecoderPreTrainedModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        embed_tokens = WrappedEmbedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+
+        self.model = SuryaADETRDecoderModel(
+            config,
+            embedder=embed_tokens,
+            static_cache=settings.LAYOUT_STATIC_CACHE,
+            max_boxes=settings.LAYOUT_MAX_BOXES
+        )
+        self.vocab_size = config.vocab_size
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    # Ignore copy
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutput]:
+        outputs = self.model(
+            input_ids=input_ids,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        return TextEncoderOutput(
+            hidden_states=outputs.last_hidden_state,
         )

surya/model/layout/encoderdecoder.py

@@ -5,7 +5,7 @@
 from transformers import PreTrainedModel, VisionEncoderDecoderConfig, PretrainedConfig
 from transformers.modeling_outputs import BaseModelOutput
 from surya.model.layout.encoder import DonutSwinLayoutModel
-from surya.model.layout.decoder import SuryaLayoutDecoder
+from surya.model.layout.decoder import SuryaLayoutDecoder, SuryaLayoutTextEncoder
 from transformers.utils import ModelOutput
 
 @dataclass
@@ -28,6 +28,7 @@ def __init__(
         config: Optional[PretrainedConfig] = None,
         encoder: Optional[PreTrainedModel] = None,
         decoder: Optional[PreTrainedModel] = None,
+        text_encoder: Optional[PreTrainedModel] = None,
     ):
         # initialize with config
         # make sure input & output embeddings is not tied
@@ -41,6 +42,10 @@ def __init__(
         if decoder is None:
             decoder = SuryaLayoutDecoder(config.decoder, attn_implementation=config._attn_implementation)
 
+        if text_encoder is None and hasattr(config, "text_encoder"):
+            text_encoder = SuryaLayoutTextEncoder(config.text_encoder, attn_implementation=config._attn_implementation)
+            self.text_encoder = text_encoder
+
         self.encoder = encoder
         self.decoder = decoder
 

surya/model/layout/model.py

@@ -1,12 +1,14 @@
 import torch
 
 from surya.model.layout.encoderdecoder import SuryaLayoutModel
-from surya.model.layout.config import SuryaLayoutConfig, SuryaLayoutDecoderConfig, DonutSwinLayoutConfig
+from surya.model.layout.config import SuryaLayoutConfig, SuryaLayoutDecoderConfig, DonutSwinLayoutConfig, \
+    SuryaLayoutTextEncoderConfig
 from surya.settings import settings
 
 
 def load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT, device=settings.TORCH_DEVICE_MODEL, dtype=settings.MODEL_DTYPE) -> SuryaLayoutModel:
     config = SuryaLayoutConfig.from_pretrained(checkpoint)
+
     decoder_config = config.decoder
     decoder = SuryaLayoutDecoderConfig(**decoder_config)
     config.decoder = decoder
@@ -15,6 +17,11 @@ def load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT, device=settings.TORC
     encoder = DonutSwinLayoutConfig(**encoder_config)
     config.encoder = encoder
 
+    if hasattr(config, "text_encoder"):
+        text_encoder_config = config.text_encoder
+        text_encoder = SuryaLayoutTextEncoderConfig(**text_encoder_config)
+        config.text_encoder = text_encoder
+
     model = SuryaLayoutModel.from_pretrained(checkpoint, config=config, torch_dtype=dtype)
     model = model.to(device)
     model = model.eval()

surya/settings.py

@@ -65,13 +65,14 @@ def TORCH_DEVICE_MODEL(self) -> str:
     RECOGNITION_ENCODER_BATCH_DIVISOR: int = 1 # Divisor for batch size in decoder
 
     # Layout
-    LAYOUT_MODEL_CHECKPOINT: str = "datalab-to/layout_order_hr4"
+    LAYOUT_MODEL_CHECKPOINT: str = "datalab-to/layout_order_te"
     LAYOUT_IMAGE_SIZE: Dict = {"height": 768, "width": 768}
     LAYOUT_BATCH_SIZE: Optional[int] = None
     LAYOUT_BENCH_DATASET_NAME: str = "vikp/publaynet_bench"
     LAYOUT_MAX_BOXES: int = 100
     COMPILE_LAYOUT: bool = False
     ORDER_BENCH_DATASET_NAME: str = "vikp/order_bench"
+    LAYOUT_MAX_DIMS: Dict = {"height": 1200, "width": 1200}
 
     # Table Rec
     TABLE_REC_MODEL_CHECKPOINT: str = "vikp/surya_tablerec"