91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
"""Measure document search Recall@K and MRR@K against a relevance set."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from search.minjust_opensearch import APP_VERSION
|
|
from search.query import search_documents
|
|
|
|
|
|
def load_queries(path: Path) -> list[dict]:
|
|
with path.open(encoding="utf-8") as source:
|
|
queries = json.load(source)
|
|
if not isinstance(queries, list) or not queries:
|
|
raise ValueError("Relevance set must be a non-empty JSON array")
|
|
|
|
seen = set()
|
|
for item in queries:
|
|
if not isinstance(item, dict) or set(item) != {
|
|
"id", "language", "query", "relevant_document_codes"
|
|
}:
|
|
raise ValueError("Each query must contain id, language, query and relevant_document_codes")
|
|
codes = item["relevant_document_codes"]
|
|
if (
|
|
not isinstance(item["id"], str)
|
|
or not item["id"].strip()
|
|
or item["id"] in seen
|
|
or not isinstance(item["language"], str)
|
|
or item["language"] not in {"ru", "ky"}
|
|
or not isinstance(item["query"], str)
|
|
or not item["query"].strip()
|
|
or not isinstance(codes, list)
|
|
or not codes
|
|
or any(not isinstance(code, str) or not code for code in codes)
|
|
or len(codes) != len(set(codes))
|
|
):
|
|
raise ValueError(f"Invalid relevance query: {item.get('id', '<unknown>')}")
|
|
seen.add(item["id"])
|
|
return queries
|
|
|
|
|
|
def search(base_url: str, index: str, item: dict, top_k: int) -> list[str]:
|
|
return search_documents(base_url, index, item["language"], item["query"], top_k)
|
|
|
|
|
|
def evaluate(queries: list[dict], base_url: str, index: str, top_k: int) -> dict:
|
|
results = []
|
|
for item in queries:
|
|
retrieved = search(base_url, index, item, top_k)
|
|
relevant = set(item["relevant_document_codes"])
|
|
matches = [rank for rank, code in enumerate(retrieved, 1) if code in relevant]
|
|
results.append({
|
|
"id": item["id"],
|
|
"language": item["language"],
|
|
"query": item["query"],
|
|
"retrieved_document_codes": retrieved,
|
|
f"recall_at_{top_k}": len(relevant.intersection(retrieved)) / len(relevant),
|
|
f"reciprocal_rank_at_{top_k}": 1 / matches[0] if matches else 0.0,
|
|
})
|
|
return {
|
|
"summary": {
|
|
"query_count": len(results),
|
|
f"recall_at_{top_k}": sum(item[f"recall_at_{top_k}"] for item in results) / len(results),
|
|
f"mrr_at_{top_k}": sum(item[f"reciprocal_rank_at_{top_k}"] for item in results) / len(results),
|
|
},
|
|
"queries": results,
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("relevance_set", type=Path)
|
|
parser.add_argument("--url", default="http://127.0.0.1:9200")
|
|
parser.add_argument("--index", default="akyldash-fragments-v1")
|
|
parser.add_argument("--top-k", type=int, default=10)
|
|
parser.add_argument("--version", action="version", version=APP_VERSION)
|
|
arguments = parser.parse_args()
|
|
if arguments.top_k <= 0:
|
|
raise SystemExit("--top-k must be greater than zero")
|
|
result = evaluate(load_queries(arguments.relevance_set), arguments.url, arguments.index, arguments.top_k)
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
print(f"Akyldash Backend v{APP_VERSION} · Frontend — not created", file=sys.stderr)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|