Benchmarking embedding models with MLflow on a managed platform
- Source
- From the RAG & Agents Engineering course
- By
- Sourav Banerjee
- Last updated
- 9 Nov 2026
The project: an embedding model bake-off on a managed platform
What you'll build
- Embeddings of one small corpus from three open models, stored side by side in one table.
- One vector search index per model, so each model is searched on its own vectors.
- A retriever per model that embeds the question with that same model and returns the top 5.
- An MLflow evaluation that grades every model's retrieval with two automatic judges.
Lesson 3 of T5.4 ended with a rule: shortlist models, then score them on your own questions. This project runs that scoring on a managed platform. Three open models from lesson 3 of T5.4's family embed the same eight documents, a managed vector search service serves each model's vectors, and MLflow's judges grade what each model retrieves for 15 queries.
Out of the notebook, the change is scale and bookkeeping. Each model gets its own index and its own MLflow run. The comparison sits in one experiment that anyone on the team can reopen, and a new candidate model is one more run.
Architecture
The pieces connect in one line, from documents to judged runs:
Figure 1: The bake-off pipeline. The 15 evaluation queries enter at the retrievers.
Each piece reuses an idea from an earlier lesson:
| Piece | What it does | Built on |
|---|---|---|
| Three embedding models | thenlper/gte-large, BAAI/bge-large-en-v1.5 and all-MiniLM-L6-v2, all with normalisation on |
Lesson 3 of T5.4, domain fit; lesson 1 of T5.4, length bias |
| One table | One row per document, with one embedding column per model | T5.2's Document and metadata |
| Three indexes | One vector search index per embedding column, created in the workspace UI | T5.5, vector stores |
| Retriever per model | Embeds the question with its own index's model and returns the top 5 | Lesson 1 of T5.4, the same model on both sides |
| Evaluation run per model | Grades every query's retrieval with a relevance judge and a groundedness judge | T5.1, baseline-first measurement |
Table 1: The project's pieces and the lessons they come from.
Following one query through the bake-off
Take the query Converting words into numbers through the MiniLM branch:
| Step | What happens |
|---|---|
| 1. Embed | The MiniLM model embeds the query, normalised to length 1 |
| 2. Search | The MiniLM index returns its 5 nearest documents with their ids and text |
| 3. Respond | The app joins the retrieved text and returns its first 300 characters as the "response" |
| 4. Judge relevance | An LLM judge rates whether each retrieved document is relevant to the query |
| 5. Judge groundedness | An LLM judge rates whether the response is supported by the retrieved documents |
| 6. Log | The scores land in the MiniLM run, with the model name logged as a parameter |
Table 2: One query through one model's branch. The GTE and BGE branches repeat the same steps on their own indexes.
The notebook recorded this result across the 15 queries:
| Model | Retrieval relevance | Retrieval groundedness |
|---|---|---|
all-MiniLM-L6-v2 |
0.36 (tied best) | 0.8 |
thenlper/gte-large |
0.36 (tied best) | 0.8 |
BAAI/bge-large-en-v1.5 |
Below 0.36 (value not given) | 0.8667 (best) |
Table 3: The recorded comparison. LLM judges can score differently on a rerun.
The smallest model tied for best relevance, and the ranking reversed on the second metric. Decide which metric your product needs before you look at results like these.
Design choices worth copying
| Choice | Why |
|---|---|
| Each retriever embeds the question with its own index's model | It rules out the silent model mismatch from lesson 1 of T5.4 |
| Normalisation on for every model | Scores compare by angle, without the length bias from lesson 1 of T5.4 |
| One run per model, with the model logged as a parameter | Runs line up side by side, one change at a time, as T5.1's baseline loop requires |
| Two judges instead of one | The ranking reversed between them, which one metric would have hidden |
| The retriever traced as its own span | The judges read the retrieved documents from the trace, not from the response text |
Table 4: Choices in the build that carry over to any model comparison.
Limits
- Eight documents and 15 queries cannot settle a model choice. They show the method. The corpus is generic machine-learning sentences, not the Larkspur handbook, and it has no hard negatives (lesson 3 of T5.4).
- Groundedness grades retrieved text here. The app returns the retrieved text itself as its response, so groundedness measures how well that excerpt matches its own sources, not an LLM's answer.
- No labelled evidence. Both scores come from LLM judges, and there is no recall@k against known evidence chunks. Run the T5.1 golden set beside the judges.
- What it needs to run: a workspace on the platform with vector search and MLflow's GenAI evaluation. The catalog and schema names are hard-coded to one user's workspace and must be changed, and the three indexes are created by hand in the UI before the evaluation cell runs.
Key Takeaways
- Give each candidate model its own index and its own evaluation run, and embed every question with the same model as its index.
- Use more than one judge: on this run the smallest model tied for best relevance while another model scored best on groundedness.
- Treat a small, unlabelled bake-off as a method check. The decision needs your own golden set, with hard negatives and labelled evidence.
Check yourself
1. Why does each retriever in the bake-off embed the question with its own index's model?
Each index holds vectors from one model, and vectors from different models live in unrelated spaces. Embedding the question with any other model would give scores without meaning and no error, the silent model mismatch from lesson 1 of T5.4.
2. MiniLM and GTE tie for best retrieval relevance at 0.36, while BGE scores best on groundedness at 0.8667. Which model do you choose?
It depends on which metric the product needs, which you should decide before reading the results. Eight documents and 15 queries also cannot settle the choice, so rerun the comparison on your own golden set before committing.
3. A colleague cites the groundedness scores as proof that BGE gives more trustworthy answers. What is wrong with that reading?
The app returns the retrieved text itself as its response, so groundedness here grades that excerpt against its own sources, not an LLM's answer. The scores say nothing about answer quality. Measure that with a real generation step and the golden set.