Upload 3 files

config.py

@@ -1,38 +1,20 @@
 from transformers import PretrainedConfig
 
 class Im2LatexTransformerConfig(PretrainedConfig):
-   """
-   Configuration for the Im2LatexTransformer model
-   """
-   model_type = "im2latex_transformer"
+   model_type = "Im2LatexTransformer"
+
+   def __init__(self, **kwargs):
+      self.vocab_size = 544
+      self.max_len = 512
+      self.d_model = 512
+      self.nhead = 8
+      self.num_layers = 6
+      self.dim_feedforward = 2048
+      self.dropout = 0.1
+      self.in_channels = 1
 
-   def __init__(
-      self,
-      vocab_size: int = 544,
-      max_len: int = 512,
-      d_model: int = 512,
-      nhead: int = 8,
-      num_decoder_layers: int = 6,
-      dim_feedforward: int = 2048,
-      in_channels: int = 1,
-      dropout: float = 0.1,
-      pad_token_id: int = 0,
-      bos_token_id: int = 1,
-      eos_token_id: int = 2,
-      **kwargs
-   ):
-      self.vocab_size = vocab_size
-      self.max_len = max_len
-      self.d_model = d_model
-      self.nhead = nhead
-      self.num_decoder_layers = num_decoder_layers
-      self.dim_feedforward = dim_feedforward
-      self.in_channels = in_channels
-      self.dropout = dropout
-      
       super().__init__(
-         pad_token_id=pad_token_id,
-         bos_token_id=bos_token_id,
-         eos_token_id=eos_token_id,
-         **kwargs
+         pad_token_id=0,
+         sos_token_id=1,
+         eos_token_id=2,
       )

image_processor.py

@@ -0,0 +1,47 @@
+from PIL import Image, ImageOps
+import numpy as np
+import torch
+from transformers import ImageProcessingMixin
+import os
+import json
+
+class Im2LatexProcessor(ImageProcessingMixin):
+   def __init__(self, image_size=(256, 256), **kwargs):
+      super().__init__(**kwargs)
+      self.image_size = image_size
+
+   def preprocess(self, image: Image.Image) -> torch.Tensor:
+      """
+      Process a PIL image and return a tensor.
+      """
+      img = image.convert("L")
+      img = ImageOps.pad(img, self.image_size, color=255)
+      arr = np.asarray(img, dtype=np.float32) / 255.0
+      arr = np.expand_dims(arr, 0)  # (1, H, W)
+      return torch.tensor(arr, dtype=torch.float32)
+
+   def __call__(self, image_path: str) -> torch.Tensor:
+      """
+      Process an image file path.
+      """
+      image = Image.open(image_path)
+      return self.preprocess(image)
+
+   def save_pretrained(self, save_directory):
+      """
+      Save processor config
+      """
+      self.image_processor_config = {
+         "image_size": self.image_size,
+      }
+      with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f:
+         json.dump(self.image_processor_config, f)
+
+   @classmethod
+   def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+      """
+      Load processor config
+      """
+      with open(os.path.join(pretrained_model_name_or_path, "preprocessor_config.json"), "r") as f:
+         config = json.load(f)
+      return cls(**config)

modeling.py

@@ -1,120 +1,202 @@
 import torch
 import torch.nn as nn
+from config import Im2LatexTransformerConfig
 from transformers import PreTrainedModel
-from config import  Im2LatexTransformerConfig
 
-# A classe CNN pode permanecer a mesma, pois é um módulo interno
 class CNN(nn.Module):
    def __init__(self, config: Im2LatexTransformerConfig):
+      """
+      Builds a CNN model
+
+      Args:
+         config (Im2LatexTransformerConfig): Configuration for the model
+      """
       super(CNN, self).__init__()
+      
       self.conv_blocks = nn.Sequential(
          nn.Conv2d(config.in_channels, 32, kernel_size=3, stride=1, padding=1),
          nn.ReLU(),
          nn.Dropout2d(p=config.dropout),
          nn.MaxPool2d(2, 2),
+
          nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
          nn.ReLU(),
          nn.Dropout2d(p=config.dropout),
          nn.MaxPool2d(2, 2),
+
          nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
          nn.ReLU(),
          nn.Dropout2d(p=config.dropout),
          nn.MaxPool2d(2, 2)
       )
+      
       self.projection = nn.Linear(128, config.d_model)
       self.dropout = nn.Dropout(config.dropout)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
+      """
+      Passes the input through the model
+
+      Args:
+          x (torch.Tensor): Input
+
+      Returns:
+          torch.Tensor: Output
+      """
+      # Garante que x tenha dimensão de batch: (B, C, H, W)
       if x.dim() == 3:
          x = x.unsqueeze(0)
-      x = self.conv_blocks(x)
+
+      # 1. Passa pelas convoluções
+      x = self.conv_blocks(x)  # -> (B, C=128, H_out, W_out)
+      
+      # 2. Prepara para o transformer
       B, C, H, W = x.shape
-      x = x.permute(0, 2, 3, 1).reshape(B, H * W, C)
-      x = self.projection(x)
+      x = x.permute(0, 2, 3, 1).reshape(B, H * W, C) # (B, S=H*W, C)
+      
+      # 3. Projeta para d_model e aplica dropout
+      x = self.projection(x) # (B, S, d_model)
       x = self.dropout(x)
+      
       return x
 
-# A classe Decoder também pode permanecer a mesma
 class Decoder(nn.Module):
    def __init__(self, config: Im2LatexTransformerConfig):
+      """
+      Builds a Transformer decoder
+
+      Args:
+         config (Im2LatexTransformerConfig): Configuration for the model
+      """
       super(Decoder, self).__init__()
       self.embedding = nn.Embedding(config.vocab_size, config.d_model)
       self.pos_embedding = nn.Embedding(config.max_len, config.d_model)
-      decoder_layer = nn.TransformerDecoderLayer(
-         config.d_model, config.nhead, config.dim_feedforward, config.dropout, batch_first=True
-      )
-      self.transformer_decoder = nn.TransformerDecoder(decoder_layer, config.num_decoder_layers)
+
+      decoder_layer = nn.TransformerDecoderLayer(config.d_model, config.nhead, config.dim_feedforward, config.dropout, batch_first=True)
+      self.transformer_decoder = nn.TransformerDecoder(decoder_layer, config.num_layers)
+
       self.output_proj = nn.Linear(config.d_model, config.vocab_size)
       self.dropout = nn.Dropout(config.dropout)
 
-   def forward(self, tokens, memory, tgt_mask=None, tgt_key_padding_mask=None):
+   def forward(self, 
+               tokens: torch.Tensor, 
+               memory: torch.Tensor, 
+               tgt_mask: torch.Tensor=None, 
+               tgt_key_padding_mask: torch.Tensor=None) -> torch.Tensor:
+      """
+      Passes the input through the decoder
+
+      Args:
+          tokens (torch.Tensor): List of tokens
+          memory (torch.Tensor): Memory
+          tgt_mask (torch.Tensor, optional): Attention mask. Defaults to None.
+          tgt_key_padding_mask (torch.Tensor, optional): Padding mask. Defaults to None.
+
+      Returns:
+          torch.Tensor: Next tokens logits
+      """
+      # tokens: (Batch, seq_len)
       batch_size, seq_len = tokens.shape
       device = tokens.device
-      token_emb = self.embedding(tokens)
-      positions = torch.arange(0, seq_len, device=device).unsqueeze(0)
-      pos_emb = self.pos_embedding(positions)
+      
+      # 1. embeddings do token + posicional
+      token_emb = self.embedding(tokens) # Shape (Batch, seq_len, d_model)
+      positions = torch.arange(0, seq_len, device=device).unsqueeze(0) # Shape (1, seq_len)
+      pos_emb = self.pos_embedding(positions) # Shape: (1, S, D)
+
+      # 2. Adiciona embeddings e aplica dropout
       x = self.dropout(token_emb + pos_emb)
+
+      # 3. Passa pelo decoder
       out = self.transformer_decoder(
          tgt=x, memory=memory, tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_key_padding_mask
       )
+
+      # 4. Passa pela projection
       logits = self.output_proj(out)
+      
       return logits
 
+# A classe Transformer wrapper também está correta.
 class Im2LatexTransformer(PreTrainedModel):
    config_class = Im2LatexTransformerConfig
 
-   def __init__(self, config: Im2LatexTransformerConfig):
-      super().__init__(config)
+   def __init__(self, config):
+      """
+      Builds a Transformer
+
+      Args:
+         config (Im2LatexTransformerConfig): Configuration for the model
+      """
+      super(Im2LatexTransformer, self).__init__(config)
       self.encoder = CNN(config)
       self.decoder = Decoder(config)
 
    def forward(self, 
                pixel_values: torch.Tensor, 
                decoder_input_ids: torch.Tensor, 
-               decoder_attention_mask: torch.Tensor=None) -> torch.Tensor:
+               decoder_padding_mask: torch.Tensor=None) -> torch.Tensor:
       """
-      A assinatura do forward é adaptada para os nomes padrão do Hugging Face.
-      - 'pixel_values' é o nome padrão para entradas de imagem.
-      - 'decoder_input_ids' é o nome padrão para os tokens do decoder.
+      Passes the input through the transformer
+
+      Args:
+         pixel_values (torch.Tensor): Input images
+         decoder_input_ids (torch.Tensor): Decoder input tokens
+         decoder_padding_mask (torch.Tensor, optional): Padding mask for the decoder. Defaults to None.
+
+      Returns:
+          torch.Tensor: Next tokens logits
       """
-      memory = self.encoder(pixel_values)
+      device = pixel_values.device
       
+      # 1. Passa pela CNN
+      memory = self.encoder(pixel_values)
+
+      # 2. Prepara o decoder
       tgt_mask = None
       if decoder_input_ids is not None:
-         device = decoder_input_ids.device
          seq_len = decoder_input_ids.size(1)
          tgt_mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.bool, device=device), diagonal=1)
 
-      logits = self.decoder(
-         tokens=decoder_input_ids, 
-         memory=memory, 
-         tgt_mask=tgt_mask, 
-         tgt_key_padding_mask=decoder_attention_mask
-      )
+      # 3. Passa pelo decoder
+      logits = self.decoder(decoder_input_ids, memory, tgt_mask, decoder_padding_mask)
       return logits
    
-   # Adicionar um método generate para facilitar a inferência
    @torch.no_grad()
-   def generate(self, pixel_values: torch.Tensor, max_length: int = 50, sos_token_id: int = 1, eos_token_id: int = 2):
-      self.eval()
+   def generate(self, pixel_values: torch.Tensor, max_length: int = 512, sos_token_id: int = 1, eos_token_id: int = 2):
+      """
+      Generates a sequence of tokens from the input images
+
+      Args:
+         pixel_values (torch.Tensor): Input images
+         max_length (int, optional): Maximum length of the generated sequence. Defaults to 512.
+         sos_token_id (int, optional): Start of sequence token ID. Defaults to 1.
+         eos_token_id (int, optional): End of sequence token ID. Defaults to 2.
+
+      Returns:
+         torch.Tensor: Generated sequence of tokens
+      """
+      self.eval()  # coloca o modelo em modo de avaliação
 
       if pixel_values.dim() == 3:
          pixel_values = pixel_values.unsqueeze(0)
+
       pixel_values = pixel_values.to(self.device)
-      memory = self.encoder(pixel_values)
-      
+
       generated_sequence = torch.tensor([[sos_token_id]], dtype=torch.long, device=self.device)
 
-      with torch.no_grad():
-         for _ in range(max_length - 1):
-            logits = self.decoder(generated_sequence, memory)
-            next_token_id = logits[0, -1, :].argmax(-1).item()
-            generated_sequence = torch.cat([
-                  generated_sequence,
-                  torch.tensor([[next_token_id]], dtype=torch.long, device=self.device)
-            ], dim=1)
-
-            if next_token_id == eos_token_id:
-                  break
-               
-      return generated_sequence
+      for _ in range(max_length):
+         logits = self(pixel_values, generated_sequence)  # forward do modelo
+         last_logits = logits[0, -1, :]  # pega a última predição
+
+         next_token_idx = last_logits.argmax(-1).item()  # greedy decoding
+
+         generated_sequence = torch.cat([
+               generated_sequence,
+               torch.tensor([[next_token_idx]], dtype=torch.long, device=self.device)
+         ], dim=1)
+
+         if next_token_idx == eos_token_id:
+               break
+
+      return generated_sequence.squeeze(0)  # remove dimensão de batch