"""Local relevance ranking for normalized evidence results."""
from __future__ import annotations
import math
import re
from collections import Counter
from collections.abc import Sequence
from dataclasses import replace
from typing import Protocol
from evidence_fetcher.models import Evidence, FeedbackLabel, RelevanceFeedback
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
ROCCHIO_ALPHA = 1.0
ROCCHIO_MORE_BETA = 3.0
ROCCHIO_RELEVANT_BETA = 3.0
ROCCHIO_LESS_GAMMA = 2.5
ROCCHIO_IRRELEVANT_GAMMA = 3.0
[docs]
class EvidenceRanker(Protocol):
"""Rank normalized evidence independently from provider retrieval."""
def rank(self, query: str, evidence: Sequence[Evidence]) -> list[Evidence]:
"""Return evidence ordered by local relevance to the query."""
[docs]
class FeedbackAwareEvidenceRanker(EvidenceRanker, Protocol):
"""Rank normalized evidence using explicit relevance feedback."""
def rerank(
self,
query: str,
evidence: Sequence[Evidence],
feedback: Sequence[RelevanceFeedback],
) -> list[Evidence]:
"""Return evidence reordered by local relevance and feedback."""
[docs]
class TfidfEvidenceRanker:
"""Small standard-library TF-IDF ranker for title and snippet text.
Feedback reranking uses a Rocchio-style query update. The original query
vector is retained with weight 1.0, MORE/RELEVANT feedback moves the query
toward marked evidence with weight 3.0, LESS moves it away with weight 2.5,
and IRRELEVANT moves it away with weight 3.0. UNSURE is retained by the
caller but has no ranking weight.
"""
def rank(self, query: str, evidence: Sequence[Evidence]) -> list[Evidence]:
query_terms = _tokenize(query)
if not evidence:
return []
if not query_terms:
return _with_ranks(evidence)
documents = [_tokenize(_document_text(item)) for item in evidence]
idf = _inverse_document_frequency(documents)
query_counts = Counter(query_terms)
scored = [
(
_score_document(query_counts, Counter(document), idf),
index,
item,
)
for index, (item, document) in enumerate(
zip(evidence, documents, strict=True)
)
]
scored.sort(key=lambda scored_item: (-scored_item[0], scored_item[1]))
return _with_ranks([item for _, _, item in scored])
def rerank(
self,
query: str,
evidence: Sequence[Evidence],
feedback: Sequence[RelevanceFeedback],
) -> list[Evidence]:
"""Rank evidence after applying explicit relevance feedback."""
if not feedback:
return self.rank(query, evidence)
if not evidence:
return []
evidence_by_url = {item.url: item for item in evidence}
unknown_urls = [
item.url for item in feedback if item.url not in evidence_by_url
]
if unknown_urls:
raise ValueError("feedback URL is not present in the evidence set")
documents = [_tokenize(_document_text(item)) for item in evidence]
idf = _inverse_document_frequency(documents)
document_vectors = [
_document_vector(Counter(document), idf) for document in documents
]
vector_by_url = {
item.url: vector
for item, vector in zip(evidence, document_vectors, strict=True)
}
query_vector = _query_vector(Counter(_tokenize(query)), idf)
revised_query_vector = _apply_feedback(query_vector, vector_by_url, feedback)
scored = [
(
_dot_product(revised_query_vector, document_vector),
index,
item,
)
for index, (item, document_vector) in enumerate(
zip(evidence, document_vectors, strict=True)
)
]
scored.sort(key=lambda scored_item: (-scored_item[0], scored_item[1]))
return _with_ranks([item for _, _, item in scored])
def _document_text(evidence: Evidence) -> str:
return f"{evidence.title} {evidence.title} {evidence.snippet}"
def _tokenize(text: str) -> list[str]:
return [match.group(0).casefold() for match in _TOKEN_RE.finditer(text)]
def _inverse_document_frequency(documents: Sequence[Sequence[str]]) -> dict[str, float]:
document_count = len(documents)
document_frequency: Counter[str] = Counter()
for document in documents:
document_frequency.update(set(document))
return {
term: math.log((1 + document_count) / (1 + frequency)) + 1
for term, frequency in document_frequency.items()
}
def _score_document(
query_counts: Counter[str],
document_counts: Counter[str],
idf: dict[str, float],
) -> float:
if not document_counts:
return 0.0
document_length = sum(document_counts.values())
score = 0.0
for term, query_count in query_counts.items():
term_frequency = document_counts[term] / document_length
score += query_count * term_frequency * idf.get(term, 0.0)
return score
def _query_vector(
query_counts: Counter[str],
idf: dict[str, float],
) -> dict[str, float]:
return {
term: count * idf.get(term, 1.0)
for term, count in query_counts.items()
if count > 0
}
def _document_vector(
document_counts: Counter[str],
idf: dict[str, float],
) -> dict[str, float]:
document_length = sum(document_counts.values())
if document_length == 0:
return {}
return {
term: count / document_length * idf[term]
for term, count in document_counts.items()
}
def _apply_feedback(
query_vector: dict[str, float],
document_vectors: dict[str, dict[str, float]],
feedback: Sequence[RelevanceFeedback],
) -> dict[str, float]:
revised = {
term: ROCCHIO_ALPHA * weight for term, weight in query_vector.items()
}
for item in feedback:
coefficient = _feedback_coefficient(item.label)
if coefficient == 0.0:
continue
document_vector = document_vectors[item.url]
for term, weight in document_vector.items():
revised[term] = revised.get(term, 0.0) + coefficient * weight
return revised
def _feedback_coefficient(label: FeedbackLabel) -> float:
match label:
case FeedbackLabel.MORE:
return ROCCHIO_MORE_BETA
case FeedbackLabel.RELEVANT:
return ROCCHIO_RELEVANT_BETA
case FeedbackLabel.LESS:
return -ROCCHIO_LESS_GAMMA
case FeedbackLabel.IRRELEVANT:
return -ROCCHIO_IRRELEVANT_GAMMA
case FeedbackLabel.UNSURE:
return 0.0
raise AssertionError(f"unsupported feedback label: {label}")
def _dot_product(left: dict[str, float], right: dict[str, float]) -> float:
if len(left) > len(right):
left, right = right, left
return sum(weight * right.get(term, 0.0) for term, weight in left.items())
def _with_ranks(evidence: Sequence[Evidence]) -> list[Evidence]:
return [replace(item, rank=index) for index, item in enumerate(evidence, start=1)]