Implementation:FlagOpen FlagEmbedding LLARA Finetune Data

Knowledge Sources	FlagOpen_FlagEmbedding
Domains	LLM Embedding, Fine-tuning, Data Loading
Last Updated	2026-02-09 00:00 GMT

Overview

Dataset and collator for fine-tuning LLARA embeddings with contrastive learning on query-passage pairs.

Description

This module provides data loading infrastructure for fine-tuning LLARA (LLM as Retrieval Adapter) embeddings. It implements TrainDatasetForEmbedding that loads query-positive-negative triplets, applies special token formatting with summarization and prediction suffixes, handles variable-length negatives through sampling, and supports loading from single files or directories. The EmbedCollator handles batch padding, optional sub-batch processing for memory efficiency, and separate query and passage tokenization. The dataset uses special tokens like <s1>-<s16> for embedding extraction points.

Usage

Use this module when fine-tuning LLARA or other LLM-based embedding models, preparing contrastive learning data with query-passage pairs, and handling datasets with variable numbers of negatives per query. The module is designed for efficient fine-tuning with large language models adapted for retrieval tasks.

Code Reference

Source Location

Repository: FlagOpen_FlagEmbedding
File: research/LLARA/finetune/data.py
Lines: 1-161

Signature

class TrainDatasetForEmbedding(Dataset):
    def __init__(
        self,
        args: DataArguments,
        tokenizer: PreTrainedTokenizer
    ):
        pass

    def __getitem__(self, item) -> Tuple[BatchEncoding, List[BatchEncoding]]:
        """Returns (query_inputs, list_of_passage_inputs)"""

@dataclass
class EmbedCollator(DataCollatorForSeq2Seq):
    query_max_len: int = 32
    passage_max_len: int = 128
    sub_batch_size: int = -1

    def __call__(self, features, return_tensors='pt'):
        """Returns dict with 'query' and 'passage' keys"""

Import

from data import TrainDatasetForEmbedding, EmbedCollator

I/O Contract

Inputs

Name	Type	Required	Description
train_data	str	Yes	Path to training JSONL file or directory
tokenizer	PreTrainedTokenizer	Yes	Tokenizer for the LLM
query_max_len	int	Yes	Maximum query length
passage_max_len	int	Yes	Maximum passage length
train_group_size	int	No	Number of passages per query (1 pos + N-1 neg)
max_example_num_per_dataset	int	No	Max examples per file
cache_path	str	No	Cache directory for datasets

Outputs

Name	Type	Description
query	BatchEncoding or List[BatchEncoding]	Tokenized queries
passage	BatchEncoding or List[BatchEncoding]	Tokenized passages (pos + negs)

Usage Examples

# Example 1: Create dataset and dataloader
from transformers import AutoTokenizer
from data import TrainDatasetForEmbedding, EmbedCollator
from torch.utils.data import DataLoader
from dataclasses import dataclass

@dataclass
class DataArguments:
    train_data: str = "./train_data"
    query_max_len: int = 32
    passage_max_len: int = 128
    train_group_size: int = 8
    max_example_num_per_dataset: int = 100000
    cache_path: str = ".cache"

# Initialize
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
args = DataArguments()

dataset = TrainDatasetForEmbedding(args=args, tokenizer=tokenizer)
collator = EmbedCollator(
    tokenizer=tokenizer,
    query_max_len=args.query_max_len,
    passage_max_len=args.passage_max_len
)

dataloader = DataLoader(
    dataset,
    batch_size=4,
    collate_fn=collator,
    shuffle=True
)

# Iterate
for batch in dataloader:
    queries = batch['query']
    passages = batch['passage']
    print(f"Query shape: {queries['input_ids'].shape}")
    print(f"Passage shape: {passages['input_ids'].shape}")
    break

# Example 2: Data format
# Training data should be JSONL with format:
# {"query": "What is ML?", "pos": ["Machine learning is..."], "neg": ["Deep learning...", "AI is..."]}

# Example 3: With sub-batching for memory efficiency
collator_sub = EmbedCollator(
    tokenizer=tokenizer,
    query_max_len=32,
    passage_max_len=128,
    sub_batch_size=32  # Process in sub-batches
)

# This will return lists of sub-batches instead of single batch
dataloader_sub = DataLoader(dataset, batch_size=4, collate_fn=collator_sub)

for batch in dataloader_sub:
    # batch['query'] and batch['passage'] are now lists of sub-batches
    if isinstance(batch['query'], list):
        print(f"Number of sub-batches: {len(batch['query'])}")
        for sub_batch in batch['query']:
            print(f"Sub-batch shape: {sub_batch['input_ids'].shape}")

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment