Denoising Objectives › BART-Style Combined Denoising

class BARTCorruptor:
    """
    Apply BART-style denoising corruption to text.

    Combines multiple corruption strategies:
    - Text infilling (spans replaced with single mask)
    - Sentence permutation (reorder sentences)
    """

    def __init__(
        self,
        mask_prob=0.30,
        mean_span_length=3,
        permute_sentences=True,
        mask_token="[MASK]",
    ):
        self.mask_prob = mask_prob
        self.mean_span_length = mean_span_length
        self.permute_sentences = permute_sentences
        self.mask_token = mask_token

    def _split_sentences(self, tokens):
        """Split token list into sentences (simple period-based)."""
        sentences = []
        current = []

        for token in tokens:
            current.append(token)
            if (
                token.endswith(".")
                or token.endswith("!")
                or token.endswith("?")
            ):
                sentences.append(current)
                current = []

        if current:
            sentences.append(current)

        return sentences

    def _infill_tokens(self, tokens):
        """Apply text infilling to token list."""
        n = len(tokens)
        if n == 0:
            return tokens

        num_to_mask = max(1, int(n * self.mask_prob))

        # Sample span lengths
        p = 1 / self.mean_span_length
        mask = np.zeros(n, dtype=bool)
        span_positions = []
        total_masked = 0

        attempts = 0
        while total_masked < num_to_mask and attempts < 100:
            attempts += 1
            span_length = np.random.geometric(p)
            span_length = min(span_length, num_to_mask - total_masked, n)

            valid_starts = [
                i
                for i in range(n - span_length + 1)
                if not any(mask[i : i + span_length])
            ]

            if not valid_starts:
                continue

            start = np.random.choice(valid_starts)
            mask[start : start + span_length] = True
            span_positions.append((start, start + span_length))
            total_masked += span_length

        span_positions.sort()

        # Build result
        result = []
        last_end = 0
        for start, end in span_positions:
            result.extend(tokens[last_end:start])
            result.append(self.mask_token)
            last_end = end
        result.extend(tokens[last_end:])

        return result

    def corrupt(self, tokens):
        """
        Apply BART-style corruption.

        Args:
            tokens: List of string tokens

        Returns:
            dict with 'input' (corrupted) and 'target' (original) token lists
        """
        working_tokens = list(tokens)

        # Step 1: Sentence permutation (if enabled and multiple sentences)
        if self.permute_sentences:
            sentences = self._split_sentences(working_tokens)
            if len(sentences) > 1:
                order = np.random.permutation(len(sentences))
                working_tokens = []
                for i in order:
                    working_tokens.extend(sentences[i])

        # Step 2: Text infilling
        corrupted = self._infill_tokens(working_tokens)

        return {"input": corrupted, "target": list(tokens)}

# Demonstrate BART corruption
document = """The transformer changed NLP. Attention is all you need.
Models became larger. Performance improved dramatically."""
tokens = document.split()

corruptor = BARTCorruptor(
    mask_prob=0.30, mean_span_length=3, permute_sentences=True
)
result = corruptor.corrupt(tokens)