Span Corruption › Implementing Span Corruption

def select_corruption_mask(
    seq_length, corruption_rate=0.15, mean_span_length=3
):
    """
    Create a boolean mask indicating which positions to corrupt.
    Spans are selected to achieve approximately the target corruption rate.
    """
    # Calculate expected number of spans
    num_tokens_to_corrupt = int(seq_length * corruption_rate)
    num_spans = max(1, num_tokens_to_corrupt // mean_span_length)

    # Sample span lengths from geometric distribution
    p = 1 / mean_span_length
    span_lengths = np.random.geometric(p, size=num_spans)

    # Clip to ensure we don't exceed the corruption budget
    total_corrupted = span_lengths.sum()
    if total_corrupted > num_tokens_to_corrupt:
        # Scale down span lengths proportionally
        scale = num_tokens_to_corrupt / total_corrupted
        span_lengths = np.maximum(1, (span_lengths * scale).astype(int))

    # Randomly place spans without overlap
    mask = np.zeros(seq_length, dtype=bool)
    available_positions = list(range(seq_length))

    for length in span_lengths:
        if len(available_positions) < length:
            break
        # Choose a starting position
        valid_starts = [i for i in range(len(available_positions) - length + 1)]
        if not valid_starts:
            break
        start_idx = np.random.choice(valid_starts)

        # Mark positions as corrupted
        for offset in range(length):
            pos = available_positions[start_idx]
            mask[pos] = True
            available_positions.remove(pos)

    return mask

# Test the span selection
test_length = 50
mask = select_corruption_mask(
    test_length, corruption_rate=0.15, mean_span_length=3
)

def get_span_boundaries(mask):
    """
    Extract start and end indices for each contiguous span in the mask.
    Returns list of (start, end) tuples where end is exclusive.
    """
    spans = []
    in_span = False
    start = 0

    for i, is_corrupted in enumerate(mask):
        if is_corrupted and not in_span:
            # Starting a new span
            start = i
            in_span = True
        elif not is_corrupted and in_span:
            # Ending current span
            spans.append((start, i))
            in_span = False

    # Handle span at end of sequence
    if in_span:
        spans.append((start, len(mask)))

    return spans

# Find spans in our test mask
spans = get_span_boundaries(mask)

def corrupt_sequence(tokens, mask, sentinel_prefix="<extra_id_"):
    """
    Apply span corruption to a token sequence.

    Args:
        tokens: List of tokens
        mask: Boolean array indicating corrupted positions
        sentinel_prefix: Prefix for sentinel tokens

    Returns:
        corrupted_input: Input sequence with sentinels replacing spans
        target: Target sequence for reconstruction
    """
    spans = get_span_boundaries(mask)

    # Build corrupted input
    corrupted_input = []
    last_end = 0

    for span_idx, (start, end) in enumerate(spans):
        # Add uncorrupted tokens before this span
        corrupted_input.extend(tokens[last_end:start])
        # Add sentinel for this span
        corrupted_input.append(f"{sentinel_prefix}{span_idx}>")
        last_end = end

    # Add remaining uncorrupted tokens
    corrupted_input.extend(tokens[last_end:])

    # Build target sequence
    target = []
    for span_idx, (start, end) in enumerate(spans):
        # Add sentinel
        target.append(f"{sentinel_prefix}{span_idx}>")
        # Add original span tokens
        target.extend(tokens[start:end])

    return corrupted_input, target

# Example sentence
sentence = "The quick brown fox jumps over the lazy dog near the riverbank"
tokens = sentence.split()

# Create corruption mask for this specific example
mask = select_corruption_mask(
    len(tokens), corruption_rate=0.25, mean_span_length=2
)

# Apply corruption
corrupted_input, target = corrupt_sequence(tokens, mask)

class SpanCorruptor:
    """Applies T5-style span corruption to text sequences."""

    def __init__(
        self, corruption_rate=0.15, mean_span_length=3, num_sentinels=100
    ):
        self.corruption_rate = corruption_rate
        self.mean_span_length = mean_span_length
        self.num_sentinels = num_sentinels
        self.sentinels = [f"<extra_id_{i}>" for i in range(num_sentinels)]

    def corrupt(self, tokens):
        """
        Corrupt a token sequence using span corruption.

        Args:
            tokens: List of string tokens

        Returns:
            dict with 'input' and 'target' token lists
        """
        if len(tokens) < 2:
            return {"input": tokens, "target": []}

        # Select spans to corrupt
        mask = select_corruption_mask(
            len(tokens), self.corruption_rate, self.mean_span_length
        )

        # Get span boundaries
        spans = get_span_boundaries(mask)

        if len(spans) > self.num_sentinels:
            # Too many spans, merge some
            spans = spans[: self.num_sentinels]

        # Build sequences
        corrupted_input = []
        target = []
        last_end = 0

        for span_idx, (start, end) in enumerate(spans):
            corrupted_input.extend(tokens[last_end:start])
            corrupted_input.append(self.sentinels[span_idx])
            target.append(self.sentinels[span_idx])
            target.extend(tokens[start:end])
            last_end = end

        corrupted_input.extend(tokens[last_end:])

        return {"input": corrupted_input, "target": target}

# Test the complete pipeline
corruptor = SpanCorruptor(corruption_rate=0.15, mean_span_length=3)

text = """Natural language processing enables computers to understand
and generate human language in meaningful ways"""
tokens = text.split()

result = corruptor.corrupt(tokens)