<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>

CC BY-SA 4.0
2024-12-18T11:45:59.253Z
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.
nbro
CC BY-SA 4.0
2024-12-19T16:52:53.330Z
@nbro hi i updated the question. I use pytorch's transformer module for this. Thanks for any input!
NatalieL
<natural-language-processing><transformer><autoencoders><loss><encoder-decoder>
Why would my transformer model always predict the same pregnancy code?
NatalieL
# 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

