Denoising Objectives › Training with Denoising Objectives

class DenoisingTrainer:
    """
    Training loop for denoising language models.

    Handles batching, corruption, and loss computation.
    """

    def __init__(self, model, tokenizer, corruptor, learning_rate=1e-4):
        self.model = model
        self.tokenizer = tokenizer
        self.corruptor = corruptor
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

    def prepare_batch(self, texts):
        """
        Prepare a batch for training.

        Args:
            texts: List of text strings

        Returns:
            dict with encoder_input, decoder_input, labels tensors
        """
        batch_encoder = []
        batch_decoder = []
        batch_labels = []

        for text in texts:
            tokens = text.split()
            result = self.corruptor.corrupt(tokens)

            # Corrupted text goes to encoder
            encoder_tokens = result["input"]
            # Original text is the target for decoder
            target_tokens = result["target"]

            batch_encoder.append(" ".join(encoder_tokens))
            batch_decoder.append(" ".join(target_tokens))
            batch_labels.append(" ".join(target_tokens))

        return {
            "encoder_texts": batch_encoder,
            "decoder_texts": batch_decoder,
            "target_texts": batch_labels,
        }

    def compute_loss(self, encoder_input, decoder_input, labels):
        """
        Compute cross-entropy loss on target tokens.

        This is a placeholder showing the structure.
        Real implementation would use model forward pass.
        """
        # In practice:
        # outputs = self.model(encoder_input, decoder_input)
        # loss = F.cross_entropy(outputs.logits, labels)
        pass