BART Pre-training › Implementation: Full Pre-training Data Pipeline

from collections import namedtuple

BARTExample = namedtuple("BARTExample", ["source", "target"])


class BARTCorruptor:
    """Handles BART-style text corruption for pre-training."""

    def __init__(
        self,
        mask_ratio: float = 0.30,
        poisson_lambda: float = 3.0,
        mask_token: str = "<mask>",
        apply_sentence_permutation: bool = True,
    ):
        self.mask_ratio = mask_ratio
        self.poisson_lambda = poisson_lambda
        self.mask_token = mask_token
        self.apply_sentence_permutation = apply_sentence_permutation

    def _split_sentences(self, text: str) -> List[str]:
        """Split text into sentences."""
        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
        return [s.strip() for s in sentences if s.strip()]

    def _permute_sentences(self, sentences: List[str]) -> List[str]:
        """Randomly permute sentence order."""
        shuffled = sentences.copy()
        random.shuffle(shuffled)
        return shuffled

    def _infill_tokens(self, tokens: List[str]) -> List[str]:
        """Apply text infilling to token sequence."""
        n = len(tokens)
        if n == 0:
            return []

        target_masked = max(1, int(n * self.mask_ratio))
        spans = []
        total_masked = 0
        attempts = 0
        max_attempts = n * 10

        while total_masked < target_masked and attempts < max_attempts:
            attempts += 1
            span_length = np.random.poisson(self.poisson_lambda)
            start = random.randint(0, n - 1)
            end = min(start + span_length, n)

            overlaps = any(not (end <= s or start >= e) for s, e in spans)

            if not overlaps:
                spans.append((start, end))
                total_masked += max(1, end - start)

        spans.sort()

        corrupted = []
        prev_end = 0
        for start, end in spans:
            corrupted.extend(tokens[prev_end:start])
            corrupted.append(self.mask_token)
            prev_end = end
        corrupted.extend(tokens[prev_end:])

        return corrupted

    def corrupt(self, text: str) -> BARTExample:
        """Apply full BART corruption pipeline."""
        # Sentence permutation
        if self.apply_sentence_permutation:
            sentences = self._split_sentences(text)
            if len(sentences) > 1:
                sentences = self._permute_sentences(sentences)
            permuted_text = " ".join(sentences)
        else:
            permuted_text = text

        # Text infilling
        tokens = permuted_text.split()
        corrupted_tokens = self._infill_tokens(tokens)

        return BARTExample(source=" ".join(corrupted_tokens), target=text)

# Demonstrate the corruptor
np.random.seed(42)
random.seed(42)

corruptor = BARTCorruptor(mask_ratio=0.30, poisson_lambda=3.0)

sample_document = """Machine learning models learn patterns from data.
Neural networks are particularly effective for complex patterns.
Pre-training helps models learn general representations."""

example = corruptor.corrupt(sample_document)