Implementation:FlagOpen FlagEmbedding RetroMAE Data
| Knowledge Sources | |
|---|---|
| Domains | Machine_Learning, Natural_Language_Processing, Self_Supervised_Learning |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
Data loading and collation for RetroMAE (Retrieval-oriented Masked Auto-Encoder) pre-training with enhanced decoder masking.
Description
This module provides data utilities for RetroMAE pre-training, which uses a unique masking strategy for encoder-decoder models:
DatasetForPretraining loads text datasets from JSON/JSONL files or HuggingFace dataset directories. It handles both single files and directories of multiple dataset files, concatenating them for pre-training.
RetroMAECollator implements the specialized masking strategy:
- Encoder masking: Applies standard whole-word masking with a configurable MLM probability (default 15%)
- Decoder masking: Uses a matrix attention mask where each token can attend to different masked versions of the sequence. For each position, it randomly selects from 128 pre-generated mask patterns, ensuring the current position is always visible. This creates diverse contexts for learning bidirectional representations.
The collator prepares inputs for both encoder MLM loss and decoder reconstruction loss, with special handling to ignore [CLS] and [SEP] tokens in labels.
Usage
Use this for pre-training encoder-decoder models with the RetroMAE objective, which is specifically designed for learning high-quality representations for retrieval tasks.
Code Reference
Source Location
- Repository: FlagOpen_FlagEmbedding
- File: research/baai_general_embedding/retromae_pretrain/data.py
- Lines: 1-100
Signature
class DatasetForPretraining(torch.utils.data.Dataset):
def __init__(self, data_dir)
def __getitem__(self, item) -> str
@dataclass
class RetroMAECollator(DataCollatorForWholeWordMask):
max_seq_length: int = 512
encoder_mlm_probability: float = 0.15
decoder_mlm_probability: float = 0.15
def __call__(self, examples)
Import
from research.baai_general_embedding.retromae_pretrain.data import DatasetForPretraining, RetroMAECollator
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| data_dir | str | Yes | Path to JSON/JSONL file or directory of dataset files |
| tokenizer | PreTrainedTokenizer | Yes | Tokenizer for encoding text |
| examples | List[str] | Yes | List of text strings to collate |
| max_seq_length | int | No | Maximum sequence length (default: 512) |
| encoder_mlm_probability | float | No | Encoder masking probability (default: 0.15) |
| decoder_mlm_probability | float | No | Decoder masking probability (default: 0.15) |
Outputs
| Name | Type | Description |
|---|---|---|
| encoder_input_ids | Tensor | Masked input IDs for encoder |
| encoder_attention_mask | Tensor | Attention mask for encoder |
| encoder_labels | Tensor | Labels for encoder MLM task |
| decoder_input_ids | Tensor | Original input IDs for decoder |
| decoder_attention_mask | Tensor | Matrix attention mask [B, L, L] for decoder |
| decoder_labels | Tensor | Labels for decoder reconstruction |
Usage Examples
from transformers import AutoTokenizer
from research.baai_general_embedding.retromae_pretrain.data import DatasetForPretraining, RetroMAECollator
# Load dataset
dataset = DatasetForPretraining(data_dir="pretrain_data/")
print(len(dataset)) # Number of text examples
# Initialize collator
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
collator = RetroMAECollator(
tokenizer=tokenizer,
max_seq_length=512,
encoder_mlm_probability=0.15,
decoder_mlm_probability=0.15
)
# Collate a batch
texts = [dataset[i] for i in range(8)]
batch = collator(texts)
# Batch contains:
# - encoder_input_ids: Masked tokens for encoder
# - encoder_labels: Original tokens for MLM loss
# - decoder_input_ids: Original tokens for decoder input
# - decoder_attention_mask: [8, L, L] matrix mask with varied patterns
# - decoder_labels: Labels for reconstruction (ignoring [CLS]/[SEP])