Implementation:Run llama Llama index CrossEncoder Dataset Gen

Overview

This module provides dataset generation utilities for cross-encoder fine-tuning. It contains two main functions: one for generating synthetic queries over documents using an LLM, and another for generating labeled query-document relevance pairs for cross-encoder training. It also defines the CrossEncoderFinetuningDatasetSample dataclass used throughout the cross-encoder fine-tuning pipeline.

Source file: llama-index-finetuning/llama_index/finetuning/cross_encoders/dataset_gen.py (174 lines)

Dependencies

Dependency	Purpose
re	Splitting LLM responses into individual questions
warnings	Warning when fewer questions are generated than requested
dataclasses.dataclass	Defining the CrossEncoderFinetuningDatasetSample data class
llama_index.core.VectorStoreIndex	Building a vector index for retrieval in dataset generation
llama_index.core.get_tokenizer	Tokenizer for the text splitter
llama_index.core.llms.ChatMessage	Structured chat messages for LLM prompting
llama_index.core.llms.llm.LLM	LLM type interface
llama_index.core.node_parser.TokenTextSplitter	Splitting documents into token-bounded chunks
llama_index.core.schema.Document	Document input type
llama_index.core.schema.MetadataMode	Controls metadata inclusion when extracting text from nodes
llama_index.llms.openai.OpenAI	Default LLM (gpt-3.5-turbo-16k) when no LLM is provided
tqdm.auto.tqdm	Progress bar display

Data Class: CrossEncoderFinetuningDatasetSample

@dataclass
class CrossEncoderFinetuningDatasetSample:
    query: str
    context: str
    score: int

A simple dataclass representing a single training sample for cross-encoder fine-tuning:

Field	Type	Description
query	str	The query text
context	str	The document or passage text
score	int	Relevance label (1 for relevant, 0 for not relevant)

Default Prompts

DEFAULT_QUERY_GEN_SYSTEM_PROMPT

Instructs the LLM to act as a professor proficient in a given topic, generating a specified number of questions separated by semicolons. The template accepts {qa_topic} and {num_questions_per_chunk} format placeholders.

DEFAULT_QUERY_GEN_USER_PROMPT

Asks the LLM to read a document and generate questions, using semicolons as delimiters. The template accepts {num_questions_per_chunk} and {context} format placeholders.

DEFAULT_QUERY_DOC_RELEVANCE_PROMPT

A multi-shot prompt for query-document relevance classification. It instructs the LLM to output a single token ("Yes" or "No") indicating whether a document is relevant to a query. The prompt includes five few-shot examples covering diverse topics (tree planting, COVID vaccine, Paris capital, PPO reinforcement learning, and sentence embeddings). The template accepts {query} and {document} format placeholders.

This prompt design is adapted from the OpenAI Cookbook cross-encoder reranking example.

Function: generate_synthetic_queries_over_documents

def generate_synthetic_queries_over_documents(
    documents: List[Document],
    num_questions_per_chunk: int = 5,
    max_chunk_length: int = 3000,
    qa_topic: str = "everything",
    llm: Optional[LLM] = None,
    qa_generate_system_msg: str = DEFAULT_QUERY_GEN_SYSTEM_PROMPT,
    qa_generate_user_msg: str = DEFAULT_QUERY_GEN_USER_PROMPT,
) -> List[str]

Parameter	Type	Default	Description
documents	List[Document]	required	Source documents for question generation
num_questions_per_chunk	int	5	Number of questions to generate per document chunk
max_chunk_length	int	3000	Maximum token length for each chunk
qa_topic	str	"everything"	Topic specialization for the question generator persona
llm	Optional[LLM]	None	LLM to use; defaults to OpenAI(model="gpt-3.5-turbo-16k", temperature=0.3)
qa_generate_system_msg	str	DEFAULT_QUERY_GEN_SYSTEM_PROMPT	System prompt template
qa_generate_user_msg	str	DEFAULT_QUERY_GEN_USER_PROMPT	User prompt template

Workflow:

1. Splits documents into chunks using TokenTextSplitter with space separator, no overlap, and newline as backup separator.
2. Extracts content from each node (without metadata).
3. For each node, formats the system and user prompts, sends them to the LLM via llm.chat().
4. Splits the LLM response on semicolons or newlines using re.split(";|\n", response_content).
5. Truncates to at most num_questions_per_chunk questions per chunk.
6. Emits a warning if fewer questions were generated than requested.
7. Returns the accumulated list of all questions.

Function: generate_ce_fine_tuning_dataset

def generate_ce_fine_tuning_dataset(
    documents: List[Document],
    questions_list: List[str],
    max_chunk_length: int = 1000,
    llm: Optional[LLM] = None,
    qa_doc_relevance_prompt: str = DEFAULT_QUERY_DOC_RELEVANCE_PROMPT,
    top_k: int = 8,
) -> List[CrossEncoderFinetuningDatasetSample]

Parameter	Type	Default	Description
documents	List[Document]	required	Source documents to build the retrieval index from
questions_list	List[str]	required	List of queries (typically from generate_synthetic_queries_over_documents)
max_chunk_length	int	1000	Maximum token length per chunk
llm	Optional[LLM]	None	LLM for relevance judgment; defaults to OpenAI(model="gpt-3.5-turbo-16k", temperature=0.1, logit_bias={9642: 1, 2822: 1})
qa_doc_relevance_prompt	str	DEFAULT_QUERY_DOC_RELEVANCE_PROMPT	Relevance prompt template
top_k	int	8	Number of documents to retrieve per query for labeling

Workflow:

1. Splits documents into chunks using the same TokenTextSplitter approach.
2. Builds a VectorStoreIndex from the chunked nodes.
3. Creates a retriever with similarity_top_k=top_k.
4. For each non-empty question:
   1. Retrieves the top-k most similar nodes.
   2. For each retrieved node, formats the relevance prompt with the query and node content, then calls llm.complete().
   3. Parses the response: if "yes", creates a sample with score=1; if "no", creates a sample with score=0; otherwise, skips the pair.
5. Returns the list of labeled CrossEncoderFinetuningDatasetSample instances.

Note on logit bias: The default OpenAI LLM is configured with logit_bias={9642: 1, 2822: 1}, which increases the likelihood of the tokens for "Yes" and "No" to constrain the output to these two choices.

Typical Usage Pipeline

Documents
    |
    v
generate_synthetic_queries_over_documents()  -->  List[str] (questions)
    |
    v
generate_ce_fine_tuning_dataset()  -->  List[CrossEncoderFinetuningDatasetSample]
    |
    v
CrossEncoderFinetuneEngine (training)

See Also

• Run_llama_Llama_index_CrossEncoderFinetuneEngine -- Cross-encoder fine-tuning engine that consumes the generated datasets
• Run_llama_Llama_index_Reranker_Dataset_Gen -- Dataset generation for Cohere reranker fine-tuning