Chunking, Retrieval Evaluation and Re-ranking

What it is

Chunking splits documents into pieces small enough to embed and retrieve, but large enough to carry a complete idea. The right strategy depends on the document structure and on two limits: the embedding model's maximum input and the context window of the LLM that reads the results.

After chunking, you evaluate retrieval to see whether the right chunks come back, and you can add a re-ranker that reorders the first results so the best ones reach the prompt.

How it works

Choosing a chunking strategy

Model constraints set the upper bound. A chunk longer than the embedding model's input limit is truncated, and the cut part is never represented. For example, databricks-bge-large-en accepts 512 tokens while databricks-gte-large-en accepts 8192.

The LLM context window limits how many chunks fit in the prompt. Chunks that are too small lose context; chunks that are too large mix topics, so their embedding is vague.

Overlap repeats the end of one chunk at the start of the next, so an idea split at a boundary still appears whole in one chunk. Too much overlap creates near duplicates that crowd the results.

Advanced chunking designs

• Parent and child (small to big): embed small child chunks, return the larger parent section to the LLM.
• Contextual headers: add the document title and section heading to each chunk's text before embedding.
• Metadata filters: store fields such as product, region or date, and filter on them at query time.
• Hybrid search: combine vector similarity with keyword matching, which helps with exact terms such as error codes.

Keep each table or code block whole in one chunk where it fits; rows split from their header lose their meaning.

Evaluating retrieval

Build an evaluation set of questions with the chunks or facts that should be retrieved, then measure.

In MLflow 3, mlflow.genai.evaluate runs scorers over traced runs. The RAG judges are RetrievalRelevance (are the retrieved chunks relevant to the query), RetrievalGroundedness (is the answer supported by the retrieved chunks) and RetrievalSufficiency (do the chunks contain everything needed). RetrievalSufficiency needs ground truth such as expected facts; the other two do not.

import mlflow
from mlflow.genai.scorers import RetrievalRelevance, RetrievalSufficiency
mlflow.genai.evaluate(
    data=eval_df,
    predict_fn=my_rag_app,
    scorers=[RetrievalRelevance(), RetrievalSufficiency()],
)

These judges read the retrieved chunks from the trace, so the retriever step must be traced.

Re-ranking

A vector index compares the query embedding with each chunk embedding separately, which is fast but approximate. A cross-encoder re-ranker reads the query and a passage together and gives a more accurate relevance score.

Because that is slower, it runs only on a small candidate set: retrieve, say, 50 chunks, re-rank them, and pass the top few to the LLM.

Mosaic AI Vector Search (called AI Search in newer docs) can apply a re-ranker within a similarity_search call through its reranker parameter, a DatabricksReranker that names the columns to score. num_results still sets how many results come back.

results = index.similarity_search(
    query_text="How do I reset a locked account?",
    columns=["chunk_id", "chunk_text", "doc_path"],
    num_results=5,
    reranker=DatabricksReranker(columns_to_rerank=["chunk_text"]),
)

How the exam asks about this

• "Answers miss details split across chunk boundaries." Add overlap or chunk by section so related text stays together.
• "Short chunks match well but the LLM lacks context to answer." Use parent and child chunking: retrieve the small chunk, return its parent.
• "The right chunk is retrieved but ranked eighth, and only the top three reach the prompt." Add a re-ranker over a larger candidate set.