Synthetic Data Generation for Fine-Tuning Custom Retrieval Models using Distilabel

In the rapidly evolving landscape of artificial intelligence, high-quality training data remains a critical challenge for machine learning practitioners. Large language models and retrieval systems require extensive, nuanced datasets to perform effectively, but manually curating such datasets is time-consuming and expensive. Enter Distilabel, an innovative framework designed to generate synthetic data generation and AI feedback, enabling AI practitioners to create sophisticated, task-specific datasets with unprecedented ease and flexibility.

Table of Content

• Introduction to Distilabel
• Hands-On Implementation
• Advanced Techniques and Best Practices

Introduction to Distilabel

Distilabel emerges as a game-changing tool in the machine learning ecosystem. By leveraging large language models and intelligent pipeline architectures, it allows researchers and developers to programmatically generate high-quality training data for various tasks, including retrieval and ranking systems. The framework’s core strength lies in its ability to automate dataset creation while maintaining high standards of data quality and diversity.

Hands-on Implementation

Let’s walk through a practical example of using Distilabel to generate a synthetic dataset for a retrieval model.

Step 1 : Setting Up the Environment

First, we install the necessary dependencies, including Distilabel with Hugging Face Inference Endpoints support:

!pip install "distilabel[hf-inference-endpoints]"
!pip install "sentence-transformers~=3.0"
!pip install "distilabel[argilla]"

Step 2 : Loading required libraries

Let’s make the needed imports:

from distilabel.llms.huggingface import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps.tasks import GenerateSentencePair
from distilabel.steps import LoadDataFromHub
from sentence_transformers import SentenceTransformer, CrossEncoder
import torch
import os
from huggingface_hub import login
import argilla as rg

login(token=os.getenv("HF_TOKEN"), add_to_git_credential=True)

Step 3 : Giving Context to LLM

context = (
"""
The text is a chunk from technical Python SDK documentation for LangChain. LangChain is a framework designed to help AI engineers and domain experts build applications powered by large language models (LLMs). It supports the creation of high-quality datasets and automates workflows. The text may include explanatory prose, Python code snippets, and references to LangChain components.
"""
)

Step 4 : Configuring the Data Generation Pipeline

We’ll use Mistral-7B-Instruct as our language model and create a pipeline with two key components: retrieval pair generation and reranking pair generation:

llm = InferenceEndpointsLLM(
    model_id="mistralai/Mistral-7B-Instruct-v0.3",
    tokenizer_id="mistralai/Mistral-7B-Instruct-v0.3",
)

with Pipeline(name="generate") as pipeline:
    load_dataset = LoadDataFromHub(
        num_examples=15,
        output_mappings={"text": "anchor"},
    )

    generate_retrieval_pairs = GenerateSentencePair(
        name="generate_retrieval_pairs",
        triplet=True,
        hard_negative=True,
        action="query",
        llm=llm,
        input_batch_size=10,
        context=context,

    )

    generate_reranking_pairs = GenerateSentencePair(
        name="generate_reranking_pairs",
        triplet=True,
        hard_negative=False,
        action="semantically-similar",
        llm=llm,
        input_batch_size=10,
        context=context,
    )
    load_dataset.connect(generate_retrieval_pairs, generate_reranking_pairs)

Step 5 : Generating and Processing the Dataset

We then run the pipeline, specifying generation parameters and loading data from the LangChain documentation:

generation_kwargs = {
    "llm": {
        "generation_kwargs": {
            "temperature": 0.7,
            "max_new_tokens": 512,
        }
    }
}
distiset = pipeline.run(

    parameters={
        load_dataset.name: {
            "repo_id": "jamescalam/langchain-docs",
            "split": "train",
        },

        generate_retrieval_pairs.name: generation_kwargs,
        generate_reranking_pairs.name: generation_kwargs,
    },
    use_cache=False,
)

Input Data:

Step 6 : Pushing Dataset to HuggingFace

distiset.push_to_hub("[Your Username]/example-retrieval-reranking-dataset2")

Generated Data:

Step 7 : Embedding and Similarity Calculation

To make the dataset useful for retrieval tasks, we embed the generated text pairs and calculate their similarities:

model_id = "sentence-transformers/all-MiniLM-L12-v2"  # Hugging Face model ID

model_retrieval = SentenceTransformer(
    model_id, device="cuda" if torch.cuda.is_available() else "cpu"
)

from sklearn.metrics.pairwise import cosine_similarity

def get_embeddings(texts):
    vectors = model_retrieval.encode(texts)
    return [vector.tolist() for vector in vectors]

def get_similarities(vector_batch_a, vector_batch_b):
    similarities = []
    for vector_a, vector_b in zip(vector_batch_a, vector_batch_b):
        similarity = cosine_similarity([vector_a], [vector_b])[0][0]
        similarities.append(similarity)
    return similarities

def format_data_retriever(batch):# -> Any:

    batch["anchor-vector"] = get_embeddings(batch["anchor"])
    batch["positive-vector"] = get_embeddings(batch["positive"])
    batch["negative-vector"] = get_embeddings(batch["negative"])    
    batch["similarity-positive-negative"] = get_similarities(batch["positive-vector"], batch["negative-vector"])
    batch["similarity-anchor-positive"] = get_similarities(batch["anchor-vector"], batch["positive-vector"])
    batch["similarity-anchor-negative"] = get_similarities(batch["anchor-vector"], batch["negative-vector"])
    return batch

dataset_generate_retrieval_pairs = distiset["generate_retrieval_pairs"]["train"].map(format_data_retriever, batched=True, batch_size=250)

Advanced Techniques and Best Practices

Advanced synthetic data generation transforms machine learning workflows by leveraging tools like Distiset, Argilla, and LLM-serving platforms. Distiset manages data efficiently, while Argilla integrates feedback for continuous refinement. Using file systems to pass data batches and caching to recover executions optimizes resource use. CLI-driven exploration enables rapid pipeline iteration, and structured data generation ensures precise output formats. Multi-model ensemble techniques and intelligent prompt engineering further enrich data quality. These strategies enable nuanced data augmentation, addressing critical challenges in training robust models across diverse domains efficiently and effectively.

Final Words

Distilabel represents more than just a tool—it’s a paradigm shift in how we approach machine learning dataset creation. By combining large language models, intelligent generation strategies, and flexible pipelines, practitioners can now generate custom, high-quality datasets tailored to specific domains and tasks. As AI continues to evolve, synthetic data generation will play an increasingly crucial role in training more specialized, accurate, and context-aware models. Frameworks like Distilabel are not just simplifying the process; they’re expanding the horizons of what’s possible in machine learning. The journey of creating better AI starts with better data, and Distilabel is leading the way.

References

• Distilabel’s Official Site
• Distilabel’s Github Repository