scripts/xkb_mmr.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""Maximal Marginal Relevance selection for recall results.
Ported from Memmy (MemTensor/memmy-agent, MIT licence) —
``Memory/src/service/retrieval/retrieval-service.ts``, functions
``mmrRecallHits`` and ``recallTextSimilarity``.
Why this is here: wiki recall kept the single highest-scoring section per
topic page. That rule is blunt in both directions — it discards a second
section even when it says something entirely different, and it lets two
near-identical sections through as long as they sit on different pages.
Once governance began giving every promoted claim its own section, a single
page could hold hundreds, and the rule was throwing away nearly all of a
page's usable content on every query.
MMR replaces "one per page" with "keep it if it adds something new": each
pick maximises ``lambda * relevance - (1 - lambda) * redundancy`` measured
against what has already been selected.
Two deliberate departures from the original:
* Memmy skips the redundancy term between two UserMemory hits, because two
things the user said are never redundant with each other. XKB has no
equivalent layer, so that special case is dropped rather than guessed at.
* Scores are normalised here. Memmy's channels already produce comparable
relevance; XKB's callers do not — wiki recall counts token hits (6 per
title match, 2 per body match, unbounded), while vector recall produces
cosine in 0..1. Feeding a count of 60 into ``0.7 * score - 0.3 * redundancy``
makes the redundancy term irrelevant and MMR a no-op. Mixing score scales
is a mistake this repository has made three times, so the normalisation is
part of the function rather than left to the caller to remember.
"""
from __future__ import annotations
import sys
from pathlib import Path
from typing import Any, Callable, Sequence
sys.path.insert(0, str(Path(__file__).resolve().parent))
import xkb_text
DEFAULT_LAMBDA = 0.7
# Redundancy is measured by xkb_text, which owns every way this project
# compares two pieces of text and documents why there is more than one.
text_terms = xkb_text.terms
text_similarity = xkb_text.similarity
def _normalised(scores: Sequence[float]) -> list[float]:
"""Map scores onto 0..1 so the redundancy term can actually compete.
Anchored at zero rather than at the lowest candidate. Min-max looks like
the obvious choice and is wrong here: it forces the weakest candidate to
0 however close it actually was, so three results scoring 60, 58 and 56
become 1.0, 0.5 and 0.0 and the redundancy term is swamped again — the
same failure the normalisation exists to prevent, one step further along.
Anchoring at zero keeps them at 1.0, 0.97 and 0.93, which is what they
are. The floor moves only when a penalty has pushed a score below zero.
"""
if not scores:
return []
low = min(0.0, min(scores))
high = max(scores)
span = high - low
if span <= 0:
return [1.0] * len(scores)
return [(score - low) / span for score in scores]
def mmr_select(
items: Sequence[Any],
limit: int,
*,
text: Callable[[Any], str],
score: Callable[[Any], float] = lambda item: float(item["score"]),
lambda_: float = DEFAULT_LAMBDA,
) -> list[Any]:
"""Pick up to ``limit`` items, trading relevance against redundancy.
``lambda_`` of 1.0 is pure relevance — the original ranking. 0.0 is pure
diversity. The default follows Memmy's 0.7.
"""
if limit <= 0 or not items:
return []
lambda_ = min(1.0, max(0.0, lambda_))
pool = list(items)
relevance = dict(zip(map(id, pool), _normalised([float(score(i)) for i in pool])))
texts = {id(item): text(item) for item in pool}
selected: list[Any] = []
while pool and len(selected) < limit:
best_index = 0
best_value = float("-inf")
for index, candidate in enumerate(pool):
redundancy = max(
(text_similarity(texts[id(candidate)], texts[id(prior)]) for prior in selected),
default=0.0,
)
value = lambda_ * relevance[id(candidate)] - (1.0 - lambda_) * redundancy
if value > best_value:
best_value = value
best_index = index
selected.append(pool.pop(best_index))
return selected