denoising with sentence permutation, and language sampling
Author: thomas-happify
Pull request: no
Labels: name: Feature request

# 🚀 Feature request

<!-- A clear and concise description of the feature proposal.
     Please provide a link to the paper and code in case they exist. -->

## Motivation

When training or fine tuning models, data collator provided in huggingface isn't enough.
For example, if we want to further pretrain `mBART` or `XLM-R`, where language sampling or sentence permutation are needed, which is hard to do with huggingface datasets API since it loads all language datasets at first.   


Thanks!


Comments:
```

import random
from typing import Dict, List, Optional, Union
import torch
from transformers import PreTrainedTokenizer
from datasets import Dataset, load_dataset

class CustomDataCollator:
    """
    Custom data collator for multilingual pretraining (e.g., mBART, XLM-R).
    Supports language sampling, sentence permutation, and lazy dataset loading.
    """
    def __init__(
        self,
        tokenizer: PreTrainedTokenizer,
        language_datasets: Dict[str, str],  # e.g., {"en": "path/to/en_dataset", "fr": "path/to/fr_dataset"}
        language_sampling_probs: Dict[str, float],  # e.g., {"en": 0.7, "fr": 0.3}
        max_length: int = 512,
        batch_size: int = 8,
        permute_sentences: bool = True,
    ):
        self.tokenizer = tokenizer
        self.language_datasets = language_datasets
        self.language_sampling_probs = language_sampling_probs
        self.max_length = max_length
        self.batch_size = batch_size
        self.permute_sentences = permute_sentences
        self.loaded_datasets = {}  # Cache for lazy-loaded datasets

    def _load_dataset(self, lang: str) -> Dataset:
        """Lazily load a dataset for a given language."""
        if lang not in self.loaded_datasets:
            # Load dataset from disk or Hugging Face hub (modify as needed)
            self.loaded_datasets[lang] = load_dataset(
                self.language_datasets[lang], split="train", streaming=True
            )
        return self.loaded_datasets[lang]

    def _sample_language(self) -> str:
        """Sample a language based on provided probabilities."""
        return random.choices(
            list(self.language_sampling_probs.keys()),
            weights=list(self.language_sampling_probs.values()),
            k=1
        )[0]

    def _permute_sentences(self, text: str) -> str:
        """Randomly permute sentences in the text."""
        sentences = text.split(". ")
        random.shuffle(sentences)
        return ". ".join(sentences)

    def __call__(self, batch: List[Dict[str, str]]) -> Dict[str, torch.Tensor]:
        """
        Collate a batch of examples. If batch is empty, dynamically sample data.
        """
        if not batch:  # If no batch provided, dynamically sample
            batch = self._sample_batch()

        # Process the batch
        input_texts = []
        for example in batch:
            text = example["text"]
            if self.permute_sentences:
                text = self._permute_sentences(text)
            input_texts.append(text)

        # Tokenize the batch
        encodings = self.tokenizer(
            input_texts,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )

        # For masked language modeling (e.g., XLM-R), add labels
        labels = encodings["input_ids"].clone()
        # Optionally mask tokens here if needed (e.g., 15% random masking)
        # For simplicity, using input_ids as labels (modify for your task)
        
        return {
            "input_ids": encodings["input_ids"],
            "attention_mask": encodings["attention_mask"],
            "labels": labels,
        }

    def _sample_batch(self) -> List[Dict[str, str]]:
        """Sample a batch dynamically from the datasets."""
        batch = []
        for _ in range(self.batch_size):
            lang = self._sample_language()
            dataset = self._load_dataset(lang)
            # Since we're using streaming=True, we iterate to get one example
            example = next(iter(dataset))
            batch.append({"text": example["text"]})  # Assumes dataset has "text" field
        return batch

# Example usage
if __name__ == "__main__":
    from transformers import MBartTokenizer

    # Initialize tokenizer (e.g., for mBART)
    tokenizer = MBartTokenizer.from_pretrained("facebook/mbart-large-cc25")

    # Define language datasets and sampling probabilities
    language_datasets = {
        "en": "wikipedia",  # Example: Hugging Face dataset name or local path
        "fr": "wikipedia",
    }
    language_sampling_probs = {
        "en": 0.7,
        "fr": 0.3,
    }

    # Create the collator
    collator = CustomDataCollator(
        tokenizer=tokenizer,
        language_datasets=language_datasets,
        language_sampling_probs=language_sampling_probs,
        max_length=128,  # Smaller for demo
        batch_size=4,
        permute_sentences=True,
    )

    # Simulate a batch
    batch = collator([])  # Empty list triggers dynamic sampling
    print("Input IDs shape:", batch["input_ids"].shape)
    print("Attention mask shape:", batch["attention_mask"].shape)
    print("Labels shape:", batch["labels"].shape)
​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​​
```