Spaces:
Sleeping
Sleeping
| """Embedding quality metrics: cosine drift, rank agreement, and IR retrieval. | |
| Used by scripts/eval_quality.py and unit-tested in tests/test_quality_metrics.py. | |
| No model or dataset dependency so the math can run in CI. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from collections.abc import Mapping, Sequence | |
| import numpy as np | |
| Qrels = Mapping[str, Mapping[str, int | float]] | |
| def l2_normalize(vectors: np.ndarray, eps: float = 1e-12) -> np.ndarray: | |
| """Row-wise L2 normalize. Zero rows stay zero (divided by eps).""" | |
| arr = np.asarray(vectors, dtype=np.float64) | |
| if arr.ndim == 1: | |
| arr = arr.reshape(1, -1) | |
| norms = np.linalg.norm(arr, axis=1, keepdims=True) | |
| return arr / np.maximum(norms, eps) | |
| def pairwise_cosine(left: np.ndarray, right: np.ndarray) -> np.ndarray: | |
| """Cosine similarity of matching rows. left[i] vs right[i].""" | |
| a = l2_normalize(left) | |
| b = l2_normalize(right) | |
| if a.shape != b.shape: | |
| raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}") | |
| return np.sum(a * b, axis=1) | |
| def cosine_matrix(queries: np.ndarray, documents: np.ndarray) -> np.ndarray: | |
| """Dense query-document cosine matrix, shape (n_queries, n_docs).""" | |
| return l2_normalize(queries) @ l2_normalize(documents).T | |
| def angular_error_deg(cosine: np.ndarray) -> np.ndarray: | |
| """Angle in degrees between two unit vectors given their cosine.""" | |
| return np.degrees(np.arccos(np.clip(np.asarray(cosine, dtype=np.float64), -1.0, 1.0))) | |
| def summarize_values(values: np.ndarray | Sequence[float]) -> dict[str, float]: | |
| arr = np.asarray(values, dtype=np.float64).reshape(-1) | |
| if arr.size == 0: | |
| return {"count": 0.0, "mean": 0.0, "p50": 0.0, "p05": 0.0, "min": 0.0, "max": 0.0} | |
| return { | |
| "count": float(arr.size), | |
| "mean": float(arr.mean()), | |
| "p50": float(np.quantile(arr, 0.50)), | |
| "p05": float(np.quantile(arr, 0.05)), | |
| "min": float(arr.min()), | |
| "max": float(arr.max()), | |
| } | |
| def cosine_drift_report(reference: np.ndarray, candidate: np.ndarray) -> dict[str, float]: | |
| """Per-row cosine and angular error of candidate embeddings vs a reference.""" | |
| cosine = pairwise_cosine(reference, candidate) | |
| angle = angular_error_deg(cosine) | |
| out = {f"cosine_{k}": v for k, v in summarize_values(cosine).items()} | |
| out.update({f"angle_deg_{k}": v for k, v in summarize_values(angle).items()}) | |
| return out | |
| def ranked_doc_ids(scores: np.ndarray, doc_ids: Sequence[str]) -> list[list[str]]: | |
| """For each query row, document ids sorted by descending score.""" | |
| if scores.ndim != 2: | |
| raise ValueError(f"scores must be 2-d, got {scores.shape}") | |
| if scores.shape[1] != len(doc_ids): | |
| raise ValueError(f"doc count {len(doc_ids)} != scores columns {scores.shape[1]}") | |
| order = np.argsort(-scores, axis=1, kind="stable") | |
| ids = list(doc_ids) | |
| return [[ids[j] for j in row] for row in order] | |
| def dcg_at_k(gains: Sequence[float], k: int) -> float: | |
| """TREC/BEIR DCG: sum_i (2^{rel_i} - 1) / log2(i + 1), 1-based ranks.""" | |
| if k <= 0: | |
| raise ValueError("k must be positive") | |
| total = 0.0 | |
| for rank, rel in enumerate(list(gains)[:k], start=1): | |
| total += (2.0 ** float(rel) - 1.0) / math.log2(rank + 1) | |
| return total | |
| def ndcg_at_k(ranked_ids: Sequence[str], qrel: Mapping[str, int | float], k: int) -> float: | |
| """nDCG@k for one query. Queries with no relevant docs score 0.""" | |
| gains = [float(qrel.get(doc_id, 0)) for doc_id in ranked_ids[:k]] | |
| dcg = dcg_at_k(gains, k) | |
| ideal = sorted((float(rel) for rel in qrel.values() if rel > 0), reverse=True) | |
| idcg = dcg_at_k(ideal, k) | |
| if idcg == 0.0: | |
| return 0.0 | |
| return dcg / idcg | |
| def recall_at_k(ranked_ids: Sequence[str], qrel: Mapping[str, int | float], k: int) -> float: | |
| """Fraction of relevant documents retrieved in the top k.""" | |
| relevant = {doc_id for doc_id, rel in qrel.items() if rel > 0} | |
| if not relevant: | |
| return 0.0 | |
| hit = sum(1 for doc_id in ranked_ids[:k] if doc_id in relevant) | |
| return hit / len(relevant) | |
| def mrr_at_k(ranked_ids: Sequence[str], qrel: Mapping[str, int | float], k: int) -> float: | |
| """Reciprocal rank of the first relevant document, 0 if none in the top k.""" | |
| relevant = {doc_id for doc_id, rel in qrel.items() if rel > 0} | |
| for rank, doc_id in enumerate(ranked_ids[:k], start=1): | |
| if doc_id in relevant: | |
| return 1.0 / rank | |
| return 0.0 | |
| def mean_retrieval_metrics( | |
| rankings: Mapping[str, Sequence[str]], | |
| qrels: Qrels, | |
| ks: Sequence[int] = (10, 100), | |
| ) -> dict[str, float]: | |
| """Macro-average nDCG/recall/MRR over queries that have at least one qrel.""" | |
| scores: dict[str, list[float]] = {} | |
| for k in ks: | |
| scores[f"ndcg@{k}"] = [] | |
| scores[f"recall@{k}"] = [] | |
| scores[f"mrr@{k}"] = [] | |
| n_scored = 0 | |
| for query_id, qrel in qrels.items(): | |
| if not any(rel > 0 for rel in qrel.values()): | |
| continue | |
| ranked = rankings.get(query_id) | |
| if ranked is None: | |
| continue | |
| n_scored += 1 | |
| for k in ks: | |
| scores[f"ndcg@{k}"].append(ndcg_at_k(ranked, qrel, k)) | |
| scores[f"recall@{k}"].append(recall_at_k(ranked, qrel, k)) | |
| scores[f"mrr@{k}"].append(mrr_at_k(ranked, qrel, k)) | |
| out: dict[str, float] = {"n_queries": float(n_scored)} | |
| for name, values in scores.items(): | |
| out[name] = float(sum(values) / len(values)) if values else 0.0 | |
| return out | |
| def topk_overlap(left: Sequence[str], right: Sequence[str], k: int) -> float: | |
| """Jaccard-style overlap of two top-k lists: |A ∩ B| / k.""" | |
| if k <= 0: | |
| raise ValueError("k must be positive") | |
| return len(set(left[:k]) & set(right[:k])) / k | |
| def mean_rank_agreement( | |
| reference_rankings: Mapping[str, Sequence[str]], | |
| candidate_rankings: Mapping[str, Sequence[str]], | |
| ks: Sequence[int] = (1, 10), | |
| ) -> dict[str, float]: | |
| """How often the candidate ranking preserves the reference top-k.""" | |
| shared = [qid for qid in reference_rankings if qid in candidate_rankings] | |
| out: dict[str, float] = {"n_queries": float(len(shared))} | |
| for k in ks: | |
| if not shared: | |
| out[f"top{k}_overlap"] = 0.0 | |
| continue | |
| overlaps = [ | |
| topk_overlap(reference_rankings[qid], candidate_rankings[qid], k) for qid in shared | |
| ] | |
| out[f"top{k}_overlap"] = float(sum(overlaps) / len(overlaps)) | |
| return out | |