RoBERTa › Implementing RoBERTa-style Training

from typing import List, Tuple

import torch


class RoBERTaDataCollator:
    """
    Data collator for RoBERTa-style pretraining.

    Key differences from BERT:
    1. No NSP - just MLM on contiguous text
    2. Dynamic masking applied fresh each batch
    3. Full sentences without artificial sentence boundaries
    """

    def __init__(
        self,
        vocab_size: int,
        mask_token_id: int,
        pad_token_id: int,
        special_token_ids: List[int],
        mask_prob: float = 0.15,
    ):
        self.vocab_size = vocab_size
        self.mask_token_id = mask_token_id
        self.pad_token_id = pad_token_id
        self.special_token_ids = set(special_token_ids)
        self.mask_prob = mask_prob

    def __call__(
        self, token_ids: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Apply dynamic masking to a batch of sequences.

        Args:
            token_ids: Shape (batch_size, seq_len)

        Returns:
            masked_ids: Input with masking applied
            labels: Original tokens at masked positions, -100 elsewhere
        """
        batch_size, seq_len = token_ids.shape
        labels = token_ids.clone()
        masked_ids = token_ids.clone()

        # Create mask for special tokens (CLS, SEP, PAD) - never mask these
        special_mask = torch.zeros_like(token_ids, dtype=torch.bool)
        for special_id in self.special_token_ids:
            special_mask |= token_ids == special_id

        # Probability matrix - 0 for special tokens, mask_prob elsewhere
        probability_matrix = torch.full(token_ids.shape, self.mask_prob)
        probability_matrix[special_mask] = 0.0

        # Sample positions to mask
        masked_indices = torch.bernoulli(probability_matrix).bool()

        # Set labels: -100 for non-masked positions (ignored in loss)
        labels[~masked_indices] = -100

        # 80% of masked tokens -> [MASK]
        indices_replaced = (
            torch.bernoulli(torch.full(token_ids.shape, 0.8)).bool()
            & masked_indices
        )
        masked_ids[indices_replaced] = self.mask_token_id

        # 10% of masked tokens -> random token
        indices_random = (
            torch.bernoulli(torch.full(token_ids.shape, 0.5)).bool()
            & masked_indices
            & ~indices_replaced
        )
        random_tokens = torch.randint(
            self.vocab_size, token_ids.shape, dtype=token_ids.dtype
        )
        masked_ids[indices_random] = random_tokens[indices_random]

        # Remaining 10% stay unchanged
        return masked_ids, labels

# Demonstrate the collator
collator = RoBERTaDataCollator(
    vocab_size=50265,  # RoBERTa vocab size
    mask_token_id=50264,  # <mask> token
    pad_token_id=1,  # <pad> token
    special_token_ids=[0, 1, 2],  # <s>, </s>, <pad>
    mask_prob=0.15,
)

# Sample batch of token IDs
torch.manual_seed(42)
sample_batch = torch.randint(100, 1000, (4, 20))  # 4 sequences, 20 tokens each

# Apply dynamic masking
masked_batch, labels = collator(sample_batch)

class FullSentenceDataset:
    """
    Dataset that packs contiguous text into sequences without NSP.

    Unlike BERT which uses sentence pairs, RoBERTa packs as much
    contiguous text as possible into each sequence.
    """

    def __init__(
        self,
        documents: List[List[int]],  # List of tokenized documents
        max_seq_len: int = 512,
        cls_token_id: int = 0,
        sep_token_id: int = 2,
    ):
        self.max_seq_len = max_seq_len
        self.cls_token_id = cls_token_id
        self.sep_token_id = sep_token_id

        # Flatten documents into sequences of max_seq_len
        self.examples = self._create_examples(documents)

    def _create_examples(
        self, documents: List[List[int]]
    ) -> List[torch.Tensor]:
        """Pack documents into fixed-length sequences."""
        examples = []
        current_chunk = []
        current_length = 0

        # Reserve space for [CLS] and [SEP]
        target_length = self.max_seq_len - 2

        for doc in documents:
            for token_id in doc:
                current_chunk.append(token_id)
                current_length += 1

                if current_length >= target_length:
                    # Create example: [CLS] tokens [SEP]
                    example = (
                        [self.cls_token_id]
                        + current_chunk[:target_length]
                        + [self.sep_token_id]
                    )
                    examples.append(torch.tensor(example))
                    current_chunk = []
                    current_length = 0

        return examples

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        return self.examples[idx]

# Demonstrate full-sentence packing
sample_docs = [
    list(range(100, 150)),  # Document 1: 50 tokens
    list(range(200, 280)),  # Document 2: 80 tokens
    list(range(300, 400)),  # Document 3: 100 tokens
]

dataset = FullSentenceDataset(
    documents=sample_docs,
    max_seq_len=64,  # Short for demonstration
    cls_token_id=0,
    sep_token_id=2,
)