Body: <p>I'm trying to predict pregnancy codes with a basic transformer model architecture. These pregnancy codes are like following prg001, prg002 to prg030. Prg001 would be antenatal screening and prg030 would be maternal outcome of delivery.</p>
  <p>The source is codes in one year (2022 for example) and the target is the codes in the following year (2023 for example).</p>
  <p>My dataset has about 42k rows. About half (22k rows) have empty target (no code because the pregnancy ended in the previous year). I use pytorch's Transformer module.</p>
  <p>The model is basically predicting zero code no matter what sequence im testing. For example, its not predicting that if a person has prg008 in the prior year, they would have prg030 (the code that represents delivery) in the target because they would give birth in the following year.</p>
  <pre><code>#Hyperparameters
  max_seq_length = 20
  max_gen_length = 20
  learning_rate = 0.001
  weight_decay = 0.01
  dropout = 0.1
  d_model = 128
  nhead = 4
  num_layers = 3
  dim_feedforward = 512
  batch_size = 256
  num_epochs = 10
  </code></pre>
  <p>I'm using crossentropy loss. Any suggestion on how I can improve this? Most of the code is below:</p>
  <pre><code>import torch
  import torch.nn as nn
  from torch.utils.data import Dataset, DataLoader
  import numpy as np
  import math
  from typing import List, Tuple
  from torch.nn.utils.rnn import pad_sequence
  import os
  import time
  
  class PreTokenizedPRGDataset(Dataset):
      def __init__(self, preprocessed_file: str):
          &quot;&quot;&quot;
          Dataset that loads pre-tokenized data from disk
          &quot;&quot;&quot;
          with open(preprocessed_file, 'rb') as f:
              self.data = pickle.load(f)
              self.data = self.data
      
      def __len__(self):
          return len(self.data)
      
      def __getitem__(self, idx):
          item = self.data[idx]
          return (
              torch.tensor(item['src'], dtype=torch.long),
              torch.tensor(item['tgt'], dtype=torch.long),
              item['id']
          )
          
  class PRGDataset(Dataset):
      def __init__(self, sequence_list, tokenizer):
          self.tokenizer = tokenizer
          # Get special token IDs
          self.bos_token_id = tokenizer.bos_token_id
          self.eos_token_id = tokenizer.eos_token_id
          self.pad_token_id = tokenizer.pad_token_id
          
          # All sequences
          self.data = []
  
          for idx, seq in enumerate(sequence_list[:50000]):
              # Convert list of tokens to IDs
              seq_src_ids = torch.tensor(tokenizer.encode(seq[0]))
              seq_tgt_ids = torch.tensor(tokenizer.encode(seq[1]))
              indiv_entpr_id = int(seq[2])
              # create source sequence
              seq_ids = torch.cat([
                  torch.tensor([self.bos_token_id]),
                  seq_src_ids,
                  torch.tensor([self.eos_token_id])
              ])
              # Create target
              tgt_ids = torch.cat([
                  torch.tensor([self.bos_token_id]),
                  seq_tgt_ids,
                  torch.tensor([self.eos_token_id])
              ])
              # # indiv_entpr_id
              # indiv_entpr_ids = torch.cat([indiv_entpr_id], dim=0)
              self.data.append((seq_ids, tgt_ids, indiv_entpr_id))
      
      def __len__(self):
          return len(self.data)
      
      def __getitem__(self, idx):
          return self.data[idx]
          
  def collate_fn(batch):
      &quot;&quot;&quot;Custom collate function to handle variable length sequences&quot;&quot;&quot;
      src_sequences = [item[0] for item in batch]
      tgt_sequences = [item[1] for item in batch]
      indiv_entpr_id = torch.tensor([item[2] for item in batch])
      
      # Pad sequences in batch
      src_padded = pad_sequence(src_sequences, batch_first=True, padding_value=0)
      tgt_padded = pad_sequence(tgt_sequences, batch_first=True, padding_value=0)
      
      return {
          'src': src_padded,
          'tgt': tgt_padded,
          'indiv_entpr_id': indiv_entpr_id
          # 'src_lengths': torch.tensor([len(s) for s in src_sequences]),
          # 'tgt_lengths': torch.tensor([len(t) for t in tgt_sequences])
      }
  
  
  class PositionalEncoding(nn.Module):
      def __init__(self, d_model: int, dropout: float, max_len: int = 20):
          super().__init__()
          self.dropout = nn.Dropout(p=dropout)
          
          # Create positional encoding matrix
          position = torch.arange(max_len).unsqueeze(1)
          div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
          
          pe = torch.zeros(1, max_len, d_model)
          pe[0, :, 0::2] = torch.sin(position * div_term)
          pe[0, :, 1::2] = torch.cos(position * div_term)
          
          # Register as buffer (not parameter)
          self.register_buffer('pe', pe)
          
      def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          Args:
              x: Tensor, shape [batch_size, seq_len, embedding_dim]
          &quot;&quot;&quot;
          x = x + self.pe[:, :x.size(1), :]
          return self.dropout(x)
  
  
  class PRGTransformer(nn.Module):
      def __init__(self, tokenizer, d_model, nhead, num_layers, dim_feedforward, max_seq_length, max_gen_length, dropout):
          super().__init__()
          self.tokenizer = tokenizer
          self.vocab_size = len(tokenizer.get_vocab())
          self.d_model = d_model
          self.pad_token_id = tokenizer.pad_token_id
          self.bos_token_id = tokenizer.bos_token_id
          self.eos_token_id = tokenizer.eos_token_id
          self.max_gen_length = max_gen_length
  
          # Token embedding
          # ***** very important to add padding_idx so model can ignore padding *****
          self.embedding = nn.Embedding(self.vocab_size, d_model, padding_idx=self.pad_token_id)
          
          # Positional encoding
          self.pos_encoder = PositionalEncoding(d_model, dropout, max_seq_length)
          
          # Dropout
          self.dropout = nn.Dropout(dropout)
          
          # Transformer
          self.transformer = nn.Transformer(
              d_model=d_model,
              nhead=nhead,
              num_encoder_layers=num_layers,
              num_decoder_layers=num_layers,
              dim_feedforward=dim_feedforward,
              dropout=dropout,
              batch_first=True,
              norm_first=True
          )
          
          # Output layer
          self.output_layer = nn.Linear(d_model, self.vocab_size)
          
          # Initialize parameters
          for p in self.parameters():
              if p.dim() &gt; 1:
                  nn.init.xavier_uniform_(p)
      
      def create_masks(self, src: torch.Tensor, tgt: torch.Tensor) -&gt; tuple:
          src_mask = None
          tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
          
          # Create padding masks if needed
          # True is where the pad tokens are
          src_padding_mask = (src == self.pad_token_id)
          tgt_padding_mask = (tgt == self.pad_token_id)
  
          return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask
      
  
      def forward(self, src, tgt):
          src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask = self.create_masks(src, tgt)
  
          # Embed tokens and add positional encoding
          src_embedded = self.embedding(src) * math.sqrt(self.d_model)
          src_embedded = self.pos_encoder(src_embedded)
          # src_embedded = self.dropout(src_embedded)
          
          tgt_embedded = self.embedding(tgt) * math.sqrt(self.d_model)
          tgt_embedded = self.pos_encoder(tgt_embedded)
          # tgt_embedded = self.dropout(tgt_embedded)
  
          # Transformer forward pass
          output = self.transformer(
              src_embedded, 
              tgt_embedded,
              src_mask=src_mask,
              tgt_mask=tgt_mask,
              src_key_padding_mask=src_key_padding_mask,
              tgt_key_padding_mask=tgt_key_padding_mask
          )
          
          return self.output_layer(output)
      
      @torch.no_grad()
      def generate(self, src: torch.Tensor) -&gt; torch.Tensor:
          &quot;&quot;&quot;
          generate sequences until EOS token is predicted
          &quot;&quot;&quot;
          self.eval()
          device = src.device
          batch_size = src.size(0)
          
          # Start with BOS token
          decoded = torch.full((batch_size, 1), self.bos_token_id, 
                             dtype=torch.long, device=device)
          
          # Generate tokens until max length (-1 for bos token)
          for _ in range(self.max_gen_length-1):
              output = self(src, decoded)
              next_token = torch.argmax(output[:, -1:, :], dim=-1)
              decoded = torch.cat([decoded, next_token], dim=1)
          
          # Find first EOS token position for each sequence, or set to sequence length if no EOS
          eos_mask = (decoded == self.eos_token_id)
          has_eos = eos_mask.any(dim=1)
          # where there's no eos, set position to sequence length
          eos_pos = torch.where(has_eos, eos_mask.float().argmax(dim=1), torch.full_like(has_eos, decoded.size(1) - 1, dtype=torch.long))
          
          # Create a mask: 1s before EOS (inclusive), 0s after
          mask = torch.arange(decoded.size(1), device=device)[None, :] &lt;= eos_pos[:, None]
          
          # Apply mask and pad with zeros
          padded = torch.where(mask, decoded, torch.zeros_like(decoded))
  
          return padded
  
  
  def train_model(model, train_loader, valid_loader, learning_rate, num_epochs, weight_decay, device, tokenizer, save_dir='checkpoints'):
      criterion = nn.CrossEntropyLoss(ignore_index=0)
      optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
      scheduler = torch.optim.lr_scheduler.OneCycleLR(
          optimizer,
          max_lr=0.001,
          epochs=num_epochs,
          steps_per_epoch=len(train_loader)
      )
      
      train_losses = []
      valid_losses = []
      best_valid_loss = float('inf')
      patience = 3        #number of epochs to wait before stopping
      
      for epoch in range(num_epochs):
          epoch_start = time.time()
          # print('epoch ', epoch)
          model.train()
          train_loss = 0
          # correct_sequences = 0
          # total_sequences = 0
          
          for batch_idx, batch in enumerate(train_loader):
              batch_start = time.time()
              src = batch['src'].to(device)
              tgt = batch['tgt'].to(device)
  
              # ***************Create teacher forcing input**************
              # 1. Create shifted version of target (shift right by 1/remove last token)
              tgt_input = tgt[:, :-1]
              # 2. Target output: remove first token
              tgt_output = tgt[:, 1:]
              
              # zero gradients out (instead of accumulating) for each batch because the previous gradients from the previous batch have already been used.
              optimizer.zero_grad()
              output = model(src, tgt_input)
              # print(f'output {output.size()}', output)
              loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))
              loss.backward()
              
              # Gradient clipping
              torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
              
              optimizer.step()
              scheduler.step()
              
              train_loss += loss.item()
          
          train_loss /= len(train_loader)         # Average training loss for the epoch
          train_losses.append(train_loss)
  
          # validation 
          print('Validation')
          model.eval()
          valid_loss = 0
          total_code_accuracy = 0
          total_sequences = 0
          with torch.no_grad():
              for batch_idx, batch in enumerate(valid_loader):
                  src = batch['src'].to(device)
                  tgt = batch['tgt'].to(device)
                  indiv_id = batch['indiv_entpr_id'].to(device)
  
                  # ***************Create teacher forcing input with SEP token**************
                  # 1. Create shifted version of target (shift right by 1/remove last token)
                  tgt_input = tgt[:, :-1]
                  # 2. Target output: remove first token
                  tgt_output = tgt[:, 1:]
                  
                  output = model(src, tgt_input)
                  # print('output', output)
                  # Reshape output and target for loss calculation
                  loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))
                  valid_loss += loss.item()
                  
                  if batch_idx % 100 == 0:
                      # Calculate accuracy
                      predictions = model.generate(src)
  
                      for pred, target in zip(predictions, tgt):
                          # print('pred', pred)
                          # print('target', target)
                          # remove padding tokens
                          pred = pred[pred != 0]
                          target = target[target != 0]
                          
                          # Convert to sets of codes for comparison
                          pred_codes = set(pred.cpu().numpy())
                          target_codes = set(target.cpu().numpy())
  
                          accuracy = len(pred_codes.intersection(target_codes)) / len(target_codes)
                          total_code_accuracy += accuracy
                          total_sequences += 1
                      
                      print(f'Batch {batch_idx} time: {time.time() - batch_start:.2f}s')
                      print(f'Epoch {epoch+1}, Batch {batch_idx}')
                      print(f'Loss: {loss.item():.4f}')
                      print(f'Sample indiv_entpr_id:     {indiv_id[0].cpu()}')
  
                      # create tensor of tokens to exclude
                      src_mask = (src[0] != tokenizer.pad_token_id) &amp; (src[0] != tokenizer.bos_token_id) &amp; (src[0] != tokenizer.eos_token_id)
                      pred_mask = (predictions[0] != tokenizer.pad_token_id) &amp; (predictions[0] != tokenizer.bos_token_id) &amp; (predictions[0] != tokenizer.eos_token_id)
                      tgt_mask = (tgt[0] != tokenizer.pad_token_id) &amp; (tgt[0] != tokenizer.bos_token_id) &amp; (tgt[0] != tokenizer.eos_token_id)
  
                      sample_input = tokenizer.decode(src[0][src_mask])
                      sample_generated = tokenizer.decode(predictions[0][pred_mask])
                      sample_target = tokenizer.decode(tgt[0][tgt_mask])
  
                      print(f'Sample Input:     {sample_input}')
                      print(f'Sample Generated: {sample_generated}')
                      print(f'Sample Target:    {sample_target}\n')
              
              valid_loss /= len(valid_loader)     # Average validation loss for the epoch
              valid_losses.append(valid_loss)
              valid_accuracy = (total_code_accuracy / total_sequences) * 100
              
          print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}, Valid Accuracy: {valid_accuracy:.2f}%')
          # Save best model and early stopping
          if valid_loss &lt; best_valid_loss:
              best_valid_loss = valid_loss
              epochs_without_improvement = 0
              
          else:
              epochs_without_improvement +=1
              if (epochs_without_improvement &gt;= patience):
                  print(f'Early stopping at epoch {epoch+1}')
                  break
          print(f'Epoch time: {time.time() - epoch_start:.2f}s')
      # Save checkpoint when we get better validation loss
      save_checkpoint(model, optimizer, scheduler, epoch, 
                  train_losses, valid_losses, best_valid_loss, 
                  save_dir)
      return train_losses, valid_losses
  
  
  def evaluate_model(model: PRGTransformer, tokenizer, test_sequences: List[torch.Tensor], batch_size, device: torch.device):
      model.eval()
      with torch.no_grad():
          test_dataset = PRGDataset(test_sequences, tokenizer)
          test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
  
          correct_sequences = 0
          total_sequences = 0
          for batch_idx, batch in enumerate(test_loader):
              src = batch['src'].to(device)
              tgt = batch['tgt'].to(device)
              
              # Calculate accuracy
              predictions = model.generate(src)
              
              # Only compare up to actual sequence lengths
              for pred, target, source in zip(predictions, tgt, src):
              
                  print(f&quot;\nInput:     {tokenizer.convert_ids_to_tokens(source)}&quot;)
                  print(f&quot;Generated: {tokenizer.convert_ids_to_tokens(pred)}&quot;)
                  print(f&quot;Expected: {tokenizer.convert_ids_to_tokens(target)}&quot;)
                  print(f&quot;Correct?   {torch.all(pred == target).item()}&quot;)
  
                  correct_sequences += torch.all(pred == target).item()
                  total_sequences += 1
          
          test_accuracy = (correct_sequences / total_sequences)
          print(f'Test Accuracy: {test_accuracy:.2%}')
  
  
  def main():
      torch.manual_seed(42)
      device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
      print(f&quot;Using device: {device}&quot;)
      
      # Load tokenizer
      tokenizer = AutoTokenizer.from_pretrained('/path/prg_tokenizer')
      
      # Hyperparameters
      max_seq_length = 20
      max_gen_length = 20
      learning_rate = 0.001
      weight_decay = 0.01
      dropout = 0.1
      d_model = 128
      nhead = 4
      num_layers = 3
      dim_feedforward = 512
      batch_size = 256
      num_epochs = 10
  
      # Create train dataset and dataloader
      train_dataset = PreTokenizedPRGDataset(train_file)
      train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, num_workers=4,
      pin_memory=True)
  
      # Create valid dataset and dataloader
      valid_dataset = PreTokenizedPRGDataset(valid_file)
      valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
      
      # Initialize model
      model =PRGTransformer(tokenizer=tokenizer, d_model=d_model, nhead=nhead, num_layers=num_layers, dim_feedforward=dim_feedforward, max_seq_length=max_seq_length, max_gen_length=max_gen_length, dropout=dropout).to(device)
      
      # Train the model
      save_dir = '/path/model_checkpoints'
      print(&quot;Starting training...&quot;)
      train_losses, valid_losses = train_model(model, train_loader, valid_loader, learning_rate, num_epochs, weight_decay, device, tokenizer, save_dir)
  
      return train_losses, valid_losses
  
  if __name__ == &quot;__main__&quot;:
      train_losses, valid_losses = main()
  </code></pre>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-12-18T11:45:59.253Z
Comments.Text: Hi. If your source code is open, it might be a good idea to share the link to it, if it's too big, or include it in the question if it's small (20 lines say). It would also be a good idea if you share more info about the tools/frameworks you're using to train this model.
Comments.UserDisplayName: nbro
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-12-19T16:52:53.330Z
Comments.Text: @nbro hi i updated the question. I use pytorch's transformer module for this. Thanks for any input!
Comments.UserDisplayName: NatalieL
Tags: <natural-language-processing><transformer><autoencoders><loss><encoder-decoder>
Title: Why would my transformer model always predict the same pregnancy code?
UserDisplayName: NatalieL
ThreadText: # Why would my transformer model always predict the same pregnancy code?
  
  (Asked by: username_0 on 2024-12-17)
  
  I'm trying to predict pregnancy codes with a basic transformer model architecture. These pregnancy codes are like following prg001, prg002 to prg030. Prg001 would be antenatal screening and prg030 would be maternal outcome of delivery.
  
  The source is codes in one year (2022 for example) and the target is the codes in the following year (2023 for example).
  
  My dataset has about 42k rows. About half (22k rows) have empty target (no code because the pregnancy ended in the previous year). I use pytorch's Transformer module.
  
  The model is basically predicting zero code no matter what sequence im testing. For example, its not predicting that if a person has prg008 in the prior year, they would have prg030 (the code that represents delivery) in the target because they would give birth in the following year.
  
  ```
  #Hyperparameters
  max_seq_length = 20
  max_gen_length = 20
  learning_rate = 0.001
  weight_decay = 0.01
  dropout = 0.1
  d_model = 128
  nhead = 4
  num_layers = 3
  dim_feedforward = 512
  batch_size = 256
  num_epochs = 10
  
  ```
  
  I'm using crossentropy loss. Any suggestion on how I can improve this? Most of the code is below:
  
  ```
  import torch
  import torch.nn as nn
  from torch.utils.data import Dataset, DataLoader
  import numpy as np
  import math
  from typing import List, Tuple
  from torch.nn.utils.rnn import pad_sequence
  import os
  import time
  
  class PreTokenizedPRGDataset(Dataset):
      def __init__(self, preprocessed_file: str):
          """
          Dataset that loads pre-tokenized data from disk
          """
          with open(preprocessed_file, 'rb') as f:
              self.data = pickle.load(f)
              self.data = self.data
      
      def __len__(self):
          return len(self.data)
      
      def __getitem__(self, idx):
          item = self.data[idx]
          return (
              torch.tensor(item['src'], dtype=torch.long),
              torch.tensor(item['tgt'], dtype=torch.long),
              item['id']
          )
          
  class PRGDataset(Dataset):
      def __init__(self, sequence_list, tokenizer):
          self.tokenizer = tokenizer
          # Get special token IDs
          self.bos_token_id = tokenizer.bos_token_id
          self.eos_token_id = tokenizer.eos_token_id
          self.pad_token_id = tokenizer.pad_token_id
          
          # All sequences
          self.data = []
  
          for idx, seq in enumerate(sequence_list[:50000]):
              # Convert list of tokens to IDs
              seq_src_ids = torch.tensor(tokenizer.encode(seq[0]))
              seq_tgt_ids = torch.tensor(tokenizer.encode(seq[1]))
              indiv_entpr_id = int(seq[2])
              # create source sequence
              seq_ids = torch.cat([
                  torch.tensor([self.bos_token_id]),
                  seq_src_ids,
                  torch.tensor([self.eos_token_id])
              ])
              # Create target
              tgt_ids = torch.cat([
                  torch.tensor([self.bos_token_id]),
                  seq_tgt_ids,
                  torch.tensor([self.eos_token_id])
              ])
              # # indiv_entpr_id
              # indiv_entpr_ids = torch.cat([indiv_entpr_id], dim=0)
              self.data.append((seq_ids, tgt_ids, indiv_entpr_id))
      
      def __len__(self):
          return len(self.data)
      
      def __getitem__(self, idx):
          return self.data[idx]
          
  def collate_fn(batch):
      """Custom collate function to handle variable length sequences"""
      src_sequences = [item[0] for item in batch]
      tgt_sequences = [item[1] for item in batch]
      indiv_entpr_id = torch.tensor([item[2] for item in batch])
      
      # Pad sequences in batch
      src_padded = pad_sequence(src_sequences, batch_first=True, padding_value=0)
      tgt_padded = pad_sequence(tgt_sequences, batch_first=True, padding_value=0)
      
      return {
          'src': src_padded,
          'tgt': tgt_padded,
          'indiv_entpr_id': indiv_entpr_id
          # 'src_lengths': torch.tensor([len(s) for s in src_sequences]),
          # 'tgt_lengths': torch.tensor([len(t) for t in tgt_sequences])
      }
  
  class PositionalEncoding(nn.Module):
      def __init__(self, d_model: int, dropout: float, max_len: int = 20):
          super().__init__()
          self.dropout = nn.Dropout(p=dropout)
          
          # Create positional encoding matrix
          position = torch.arange(max_len).unsqueeze(1)
          div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
          
          pe = torch.zeros(1, max_len, d_model)
          pe[0, :, 0::2] = torch.sin(position * div_term)
          pe[0, :, 1::2] = torch.cos(position * div_term)
          
          # Register as buffer (not parameter)
          self.register_buffer('pe', pe)
          
      def forward(self, x: torch.Tensor) -> torch.Tensor:
          """
          Args:
              x: Tensor, shape [batch_size, seq_len, embedding_dim]
          """
          x = x + self.pe[:, :x.size(1), :]
          return self.dropout(x)
  
  class PRGTransformer(nn.Module):
      def __init__(self, tokenizer, d_model, nhead, num_layers, dim_feedforward, max_seq_length, max_gen_length, dropout):
          super().__init__()
          self.tokenizer = tokenizer
          self.vocab_size = len(tokenizer.get_vocab())
          self.d_model = d_model
          self.pad_token_id = tokenizer.pad_token_id
          self.bos_token_id = tokenizer.bos_token_id
          self.eos_token_id = tokenizer.eos_token_id
          self.max_gen_length = max_gen_length
  
          # Token embedding
          # ***** very important to add padding_idx so model can ignore padding *****
          self.embedding = nn.Embedding(self.vocab_size, d_model, padding_idx=self.pad_token_id)
          
          # Positional encoding
          self.pos_encoder = PositionalEncoding(d_model, dropout, max_seq_length)
          
          # Dropout
          self.dropout = nn.Dropout(dropout)
          
          # Transformer
          self.transformer = nn.Transformer(
              d_model=d_model,
              nhead=nhead,
              num_encoder_layers=num_layers,
              num_decoder_layers=num_layers,
              dim_feedforward=dim_feedforward,
              dropout=dropout,
              batch_first=True,
              norm_first=True
          )
          
          # Output layer
          self.output_layer = nn.Linear(d_model, self.vocab_size)
          
          # Initialize parameters
          for p in self.parameters():
              if p.dim() > 1:
                  nn.init.xavier_uniform_(p)
      
      def create_masks(self, src: torch.Tensor, tgt: torch.Tensor) -> tuple:
          src_mask = None
          tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
          
          # Create padding masks if needed
          # True is where the pad tokens are
          src_padding_mask = (src == self.pad_token_id)
          tgt_padding_mask = (tgt == self.pad_token_id)
  
          return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask
      
  
      def forward(self, src, tgt):
          src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask = self.create_masks(src, tgt)
  
          # Embed tokens and add positional encoding
          src_embedded = self.embedding(src) * math.sqrt(self.d_model)
          src_embedded = self.pos_encoder(src_embedded)
          # src_embedded = self.dropout(src_embedded)
          
          tgt_embedded = self.embedding(tgt) * math.sqrt(self.d_model)
          tgt_embedded = self.pos_encoder(tgt_embedded)
          # tgt_embedded = self.dropout(tgt_embedded)
  
          # Transformer forward pass
          output = self.transformer(
              src_embedded, 
              tgt_embedded,
              src_mask=src_mask,
              tgt_mask=tgt_mask,
              src_key_padding_mask=src_key_padding_mask,
              tgt_key_padding_mask=tgt_key_padding_mask
          )
          
          return self.output_layer(output)
      
      @torch.no_grad()
      def generate(self, src: torch.Tensor) -> torch.Tensor:
          """
          generate sequences until EOS token is predicted
          """
          self.eval()
          device = src.device
          batch_size = src.size(0)
          
          # Start with BOS token
          decoded = torch.full((batch_size, 1), self.bos_token_id, 
                             dtype=torch.long, device=device)
          
          # Generate tokens until max length (-1 for bos token)
          for _ in range(self.max_gen_length-1):
              output = self(src, decoded)
              next_token = torch.argmax(output[:, -1:, :], dim=-1)
              decoded = torch.cat([decoded, next_token], dim=1)
          
          # Find first EOS token position for each sequence, or set to sequence length if no EOS
          eos_mask = (decoded == self.eos_token_id)
          has_eos = eos_mask.any(dim=1)
          # where there's no eos, set position to sequence length
          eos_pos = torch.where(has_eos, eos_mask.float().argmax(dim=1), torch.full_like(has_eos, decoded.size(1) - 1, dtype=torch.long))
          
          # Create a mask: 1s before EOS (inclusive), 0s after
          mask = torch.arange(decoded.size(1), device=device)[None, :] <= eos_pos[:, None]
          
          # Apply mask and pad with zeros
          padded = torch.where(mask, decoded, torch.zeros_like(decoded))
  
          return padded
  
  def train_model(model, train_loader, valid_loader, learning_rate, num_epochs, weight_decay, device, tokenizer, save_dir='checkpoints'):
      criterion = nn.CrossEntropyLoss(ignore_index=0)
      optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
      scheduler = torch.optim.lr_scheduler.OneCycleLR(
          optimizer,
          max_lr=0.001,
          epochs=num_epochs,
          steps_per_epoch=len(train_loader)
      )
      
      train_losses = []
      valid_losses = []
      best_valid_loss = float('inf')
      patience = 3        #number of epochs to wait before stopping
      
      for epoch in range(num_epochs):
          epoch_start = time.time()
          # print('epoch ', epoch)
          model.train()
          train_loss = 0
          # correct_sequences = 0
          # total_sequences = 0
          
          for batch_idx, batch in enumerate(train_loader):
              batch_start = time.time()
              src = batch['src'].to(device)
              tgt = batch['tgt'].to(device)
  
              # ***************Create teacher forcing input**************
              # 1. Create shifted version of target (shift right by 1/remove last token)
              tgt_input = tgt[:, :-1]
              # 2. Target output: remove first token
              tgt_output = tgt[:, 1:]
              
              # zero gradients out (instead of accumulating) for each batch because the previous gradients from the previous batch have already been used.
              optimizer.zero_grad()
              output = model(src, tgt_input)
              # print(f'output {output.size()}', output)
              loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))
              loss.backward()
              
              # Gradient clipping
              torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
              
              optimizer.step()
              scheduler.step()
              
              train_loss += loss.item()
          
          train_loss /= len(train_loader)         # Average training loss for the epoch
          train_losses.append(train_loss)
  
          # validation 
          print('Validation')
          model.eval()
          valid_loss = 0
          total_code_accuracy = 0
          total_sequences = 0
          with torch.no_grad():
              for batch_idx, batch in enumerate(valid_loader):
                  src = batch['src'].to(device)
                  tgt = batch['tgt'].to(device)
                  indiv_id = batch['indiv_entpr_id'].to(device)
  
                  # ***************Create teacher forcing input with SEP token**************
                  # 1. Create shifted version of target (shift right by 1/remove last token)
                  tgt_input = tgt[:, :-1]
                  # 2. Target output: remove first token
                  tgt_output = tgt[:, 1:]
                  
                  output = model(src, tgt_input)
                  # print('output', output)
                  # Reshape output and target for loss calculation
                  loss = criterion(output.reshape(-1, output.size(-1)), tgt_output.reshape(-1))
                  valid_loss += loss.item()
                  
                  if batch_idx % 100 == 0:
                      # Calculate accuracy
                      predictions = model.generate(src)
  
                      for pred, target in zip(predictions, tgt):
                          # print('pred', pred)
                          # print('target', target)
                          # remove padding tokens
                          pred = pred[pred != 0]
                          target = target[target != 0]
                          
                          # Convert to sets of codes for comparison
                          pred_codes = set(pred.cpu().numpy())
                          target_codes = set(target.cpu().numpy())
  
                          accuracy = len(pred_codes.intersection(target_codes)) / len(target_codes)
                          total_code_accuracy += accuracy
                          total_sequences += 1
                      
                      print(f'Batch {batch_idx} time: {time.time() - batch_start:.2f}s')
                      print(f'Epoch {epoch+1}, Batch {batch_idx}')
                      print(f'Loss: {loss.item():.4f}')
                      print(f'Sample indiv_entpr_id:     {indiv_id[0].cpu()}')
  
                      # create tensor of tokens to exclude
                      src_mask = (src[0] != tokenizer.pad_token_id) & (src[0] != tokenizer.bos_token_id) & (src[0] != tokenizer.eos_token_id)
                      pred_mask = (predictions[0] != tokenizer.pad_token_id) & (predictions[0] != tokenizer.bos_token_id) & (predictions[0] != tokenizer.eos_token_id)
                      tgt_mask = (tgt[0] != tokenizer.pad_token_id) & (tgt[0] != tokenizer.bos_token_id) & (tgt[0] != tokenizer.eos_token_id)
  
                      sample_input = tokenizer.decode(src[0][src_mask])
                      sample_generated = tokenizer.decode(predictions[0][pred_mask])
                      sample_target = tokenizer.decode(tgt[0][tgt_mask])
  
                      print(f'Sample Input:     {sample_input}')
                      print(f'Sample Generated: {sample_generated}')
                      print(f'Sample Target:    {sample_target}\n')
              
              valid_loss /= len(valid_loader)     # Average validation loss for the epoch
              valid_losses.append(valid_loss)
              valid_accuracy = (total_code_accuracy / total_sequences) * 100
              
          print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}, Valid Accuracy: {valid_accuracy:.2f}%')
          # Save best model and early stopping
          if valid_loss < best_valid_loss:
              best_valid_loss = valid_loss
              epochs_without_improvement = 0
              
          else:
              epochs_without_improvement +=1
              if (epochs_without_improvement >= patience):
                  print(f'Early stopping at epoch {epoch+1}')
                  break
          print(f'Epoch time: {time.time() - epoch_start:.2f}s')
      # Save checkpoint when we get better validation loss
      save_checkpoint(model, optimizer, scheduler, epoch, 
                  train_losses, valid_losses, best_valid_loss, 
                  save_dir)
      return train_losses, valid_losses
  
  def evaluate_model(model: PRGTransformer, tokenizer, test_sequences: List[torch.Tensor], batch_size, device: torch.device):
      model.eval()
      with torch.no_grad():
          test_dataset = PRGDataset(test_sequences, tokenizer)
          test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
  
          correct_sequences = 0
          total_sequences = 0
          for batch_idx, batch in enumerate(test_loader):
              src = batch['src'].to(device)
              tgt = batch['tgt'].to(device)
              
              # Calculate accuracy
              predictions = model.generate(src)
              
              # Only compare up to actual sequence lengths
              for pred, target, source in zip(predictions, tgt, src):
              
                  print(f"\nInput:     {tokenizer.convert_ids_to_tokens(source)}")
                  print(f"Generated: {tokenizer.convert_ids_to_tokens(pred)}")
                  print(f"Expected: {tokenizer.convert_ids_to_tokens(target)}")
                  print(f"Correct?   {torch.all(pred == target).item()}")
  
                  correct_sequences += torch.all(pred == target).item()
                  total_sequences += 1
          
          test_accuracy = (correct_sequences / total_sequences)
          print(f'Test Accuracy: {test_accuracy:.2%}')
  
  def main():
      torch.manual_seed(42)
      device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
      print(f"Using device: {device}")
      
      # Load tokenizer
      tokenizer = AutoTokenizer.from_pretrained('/path/prg_tokenizer')
      
      # Hyperparameters
      max_seq_length = 20
      max_gen_length = 20
      learning_rate = 0.001
      weight_decay = 0.01
      dropout = 0.1
      d_model = 128
      nhead = 4
      num_layers = 3
      dim_feedforward = 512
      batch_size = 256
      num_epochs = 10
  
      # Create train dataset and dataloader
      train_dataset = PreTokenizedPRGDataset(train_file)
      train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn, num_workers=4,
      pin_memory=True)
  
      # Create valid dataset and dataloader
      valid_dataset = PreTokenizedPRGDataset(valid_file)
      valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
      
      # Initialize model
      model =PRGTransformer(tokenizer=tokenizer, d_model=d_model, nhead=nhead, num_layers=num_layers, dim_feedforward=dim_feedforward, max_seq_length=max_seq_length, max_gen_length=max_gen_length, dropout=dropout).to(device)
      
      # Train the model
      save_dir = '/path/model_checkpoints'
      print("Starting training...")
      train_losses, valid_losses = train_model(model, train_loader, valid_loader, learning_rate, num_epochs, weight_decay, device, tokenizer, save_dir)
  
      return train_losses, valid_losses
  
  if __name__ == "__main__":
      train_losses, valid_losses = main()
  
  ```
  
  **Tags:** natural-language-processing, transformer, autoencoders, loss, encoder-decoder
  
  > **Comments:**
  > **username_1**: Hi. If your source code is open, it might be a good idea to share the link to it, if it's too big, or include it in the question if it's small (20 lines say). It would also be a good idea if you share more info about the tools/frameworks you're using to train this model.  
  > **username_0**: @username_1 hi i updated the question. I use pytorch's transformer module for this. Thanks for any input!  
  
  ---
  
  ## 0 answers
  
  