Maximal Marginal Relevance (MMR) Selection Stephard

Maximal Marginal Relevance (MMR) Selection Step

hard RAG & Retrieval 15 min

NumPy allowed · plain Python earns 10 bonus points

A query retrieves context

Embedding similarity ranks the documents before generation.

Open lesson →

01query

02embed

03retrieve

04context

What is RAG?→ embed →top context

retrieval.md44%

rag-systems.md33%

agents.md24%

Cosine-like similarity scores ranked before context assembly

Implement one selection step of Maximal Marginal Relevance, the real re-ranking algorithm that balances relevance against redundancy when building a diverse retrieved set.

Task

Implement mmr_select_next(query_sims, selected, doc_sims, lambda_param). Among candidate doc ids NOT already in selected, each candidate d's MMR score is lambda_param * query_sims[d] - (1-lambda_param) * max(pairwise_sim(d,s) for s in selected) (if selected is empty, the second term is 0). Pairwise similarity between two ids is looked up in doc_sims under the key f'{a},{b}' OR f'{b},{a}' (whichever exists). Return the id of the candidate with the HIGHEST MMR score; ties broken by id ascending.

Function Signature

mmr_select_next(query_sims: dict[str, float], selected: list[str], doc_sims: dict[str, float], lambda_param: float) -> str

Examples

Example 1: first pick is pure relevance
Input: query_sims = {'a': 0.9, 'b': 0.5}
selected = []
doc_sims = {}
lambda_param = 0.5
Output: 'a'

Example 2: redundant doc penalized
Input: query_sims = {'a': 0.9, 'b': 0.85}
selected = ['a']
doc_sims = {'a,b': 0.95}
lambda_param = 0.5
Output: 'b'

Example 3: diverse doc preferred despite lower relevance
Input: query_sims = {'a': 0.9, 'b': 0.7}
selected = ['a']
doc_sims = {'a,b': 0.9}
lambda_param = 0.3
Output: 'b'

Constraints

query_sims non-empty
0 <= lambda_param <= 1

Python3Saved ✓

def mmr_select_next(query_sims, selected, doc_sims, lambda_param):

# Your implementation here

pass

Case 1: first pick is pure relevance

Input: query_sims = {'a': 0.9, 'b': 0.5}
selected = []
doc_sims = {}
lambda_param = 0.5

Expected: 'a'

Case 2: redundant doc penalized

Input: query_sims = {'a': 0.9, 'b': 0.85}
selected = ['a']
doc_sims = {'a,b': 0.95}
lambda_param = 0.5

Expected: 'b'

Case 3: diverse doc preferred despite lower relevance

Input: query_sims = {'a': 0.9, 'b': 0.7}
selected = ['a']
doc_sims = {'a,b': 0.9}
lambda_param = 0.3

Expected: 'b'