from typing import Optional, Tuple, Union import torch from torch import nn from surya.common.adetr.decoder import SuryaADETRDecoderModel, SuryaADETRDecoderPreTrainedModel from surya.table_rec.model.config import TableRecModelOutput from surya.table_rec.shaper import LabelShaper from surya.settings import settings class LabelEmbedding(nn.Module): def __init__(self, config): super().__init__() # Bboxes self.w_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.h_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.cx_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.cy_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.xskew_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.yskew_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.x1_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.y1_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.x2_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.y2_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.x3_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.y3_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.x4_embed = nn.Embedding(config.vocab_size, config.box_embed_size) self.y4_embed = nn.Embedding(config.vocab_size, config.box_embed_size) # Get indexes for passed in tensor shaper = LabelShaper() self.component_idxs = shaper.component_idx_dict() merge_count = shaper.get_box_property("merges")[1] + config.special_token_count category_count = shaper.get_box_property("category")[1] + config.special_token_count # Other box properties self.category_embed = nn.Embedding(category_count, config.property_embed_size) self.merge_embed = nn.Embedding(merge_count, config.property_embed_size) self.colspan_embed = nn.Embedding(config.vocab_size, config.property_embed_size) self.config = config def forward(self, boxes: torch.LongTensor, *args): # Need to keep *args for compatibility with common decoder boxes = boxes.to(torch.long).clamp(0, self.config.vocab_size) boxes_unbound = boxes.to(torch.long).unbind(dim=-1) cx, cy, w, h, xskew, yskew = boxes_unbound[self.component_idxs["bbox"][0]:self.component_idxs["bbox"][1]] category = boxes_unbound[self.component_idxs["category"][0]:self.component_idxs["category"][1]][0] merges = boxes_unbound[self.component_idxs["merges"][0]:self.component_idxs["merges"][1]][0] colspan = boxes_unbound[self.component_idxs["colspan"][0]:self.component_idxs["colspan"][1]][0] xskew_actual = ((xskew - self.config.bbox_size // 2) / 2).to(torch.long) yskew_actual = ((yskew - self.config.bbox_size // 2) / 2).to(torch.long) x1 = (cx - w // 2 - xskew_actual).clamp(0, self.config.bbox_size).to(torch.long) y1 = (cy - h // 2 - yskew_actual).clamp(0, self.config.bbox_size).to(torch.long) x3 = (cx + w // 2 + xskew_actual).clamp(0, self.config.bbox_size).to(torch.long) y3 = (cy + h // 2 + yskew_actual).clamp(0, self.config.bbox_size).to(torch.long) size_embeds = self.w_embed(w) + self.h_embed(h) + self.cx_embed(cx) + self.cy_embed(cy) skew_embeds = self.xskew_embed(xskew) + self.yskew_embed(yskew) corner_embeds = self.x1_embed(x1) + self.y1_embed(y1) + self.x3_embed(x3) + self.y3_embed(y3) box_embeds = size_embeds + skew_embeds + corner_embeds property_embeds = self.category_embed(category) + self.merge_embed(merges) + self.colspan_embed(colspan) # Cat bbox and property embeddings embedded = torch.cat([box_embeds, property_embeds], dim=-1) return embedded class SuryaTableRecDecoder(SuryaADETRDecoderPreTrainedModel): _tied_weights_keys = None def __init__(self, config, **kwargs): super().__init__(config) embed_tokens = LabelEmbedding(config) self.model = SuryaADETRDecoderModel( config, embedder=embed_tokens, static_cache=settings.TABLE_REC_STATIC_CACHE, max_boxes=settings.TABLE_REC_MAX_BOXES ) self.vocab_size = config.vocab_size shaper = LabelShaper() property_heads = {} for k in shaper.property_keys: _, kcount, mode = shaper.get_box_property(k) property_heads[k] = nn.Linear(config.hidden_size, kcount, bias=False) self.box_property_heads = nn.ModuleDict(property_heads) self.pre_output_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) # Initialize weights and apply final processing 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 get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings 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, prefill: bool = False, **kwargs ) -> Union[Tuple, TableRecModelOutput]: 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, prefill=prefill, ) hidden_states = self.pre_output_norm(outputs[0]) box_property_logits = {} for key in self.box_property_heads: box_property_logits[key] = self.box_property_heads[key](hidden_states) bbox_logits = nn.functional.sigmoid(box_property_logits["bbox"]) box_property_logits["bbox"] = bbox_logits return TableRecModelOutput( box_property_logits=box_property_logits, hidden_states=hidden_states, )