Implementation:FlagOpen FlagEmbedding LLARA Finetune Data
| Knowledge Sources | |
|---|---|
| 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}")