406 lines
14 KiB
Python
Executable File
406 lines
14 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import html
|
|
import json
|
|
import re
|
|
import unicodedata
|
|
from collections import Counter
|
|
from difflib import SequenceMatcher
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from ir_measures import P, RR, Recall, nDCG, Qrel, ScoredDoc, calc_aggregate
|
|
|
|
|
|
CONFIGS = [
|
|
"A0_full",
|
|
"A1_no_planner",
|
|
"A2_no_neo4j",
|
|
"A3_no_vector",
|
|
"A4_one_hop",
|
|
"A5_no_bge",
|
|
]
|
|
KS = (5, 10, 30)
|
|
|
|
|
|
def normalize(text: Any) -> str:
|
|
value = html.unescape(str(text or ""))
|
|
value = unicodedata.normalize("NFKC", value).lower()
|
|
value = re.sub(r"<[^>]+>", " ", value)
|
|
value = re.sub(r"https?://\S+", " ", value)
|
|
value = re.sub(r"[^a-z0-9\u4e00-\u9fff.%+-]+", " ", value)
|
|
return " ".join(value.split())
|
|
|
|
|
|
def tokens(text: Any) -> list[str]:
|
|
return re.findall(r"[a-z0-9.%+-]+|[\u4e00-\u9fff]", normalize(text))
|
|
|
|
|
|
def evidence_similarity(gold: str, candidate: str) -> dict[str, float | bool]:
|
|
ng = normalize(gold)
|
|
nc = normalize(candidate)
|
|
if not ng or not nc:
|
|
return {
|
|
"matched": False,
|
|
"containment": False,
|
|
"token_recall": 0.0,
|
|
"token_precision": 0.0,
|
|
"sequence_ratio": 0.0,
|
|
"match_score": 0.0,
|
|
}
|
|
containment = ng in nc or nc in ng
|
|
gt = set(tokens(gold))
|
|
ct = set(tokens(candidate))
|
|
overlap = len(gt & ct)
|
|
token_recall = overlap / len(gt) if gt else 0.0
|
|
token_precision = overlap / len(ct) if ct else 0.0
|
|
sequence_ratio = SequenceMatcher(None, ng, nc).ratio()
|
|
match_score = max(
|
|
sequence_ratio,
|
|
0.72 * token_recall + 0.28 * token_precision,
|
|
)
|
|
matched = bool(
|
|
containment
|
|
or sequence_ratio >= 0.82
|
|
or (token_recall >= 0.72 and token_precision >= 0.28)
|
|
)
|
|
return {
|
|
"matched": matched,
|
|
"containment": containment,
|
|
"token_recall": token_recall,
|
|
"token_precision": token_precision,
|
|
"sequence_ratio": sequence_ratio,
|
|
"match_score": match_score,
|
|
}
|
|
|
|
|
|
def evidence_doc_id(evidence: str) -> str:
|
|
digest = hashlib.sha1(normalize(evidence).encode("utf-8")).hexdigest()[:20]
|
|
return f"EVID::{digest}"
|
|
|
|
|
|
def paper_id(source: dict[str, Any]) -> str:
|
|
return str(
|
|
source.get("paper_id")
|
|
or source.get("pmcid")
|
|
or source.get("pmid")
|
|
or source.get("doi")
|
|
or source.get("paper_title")
|
|
or source.get("title")
|
|
or ""
|
|
).strip()
|
|
|
|
|
|
def dedupe_ranked_sources(
|
|
qid: str,
|
|
gold_source: str,
|
|
sources: list[dict[str, Any]],
|
|
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
|
ranked: list[dict[str, Any]] = []
|
|
audit: list[dict[str, Any]] = []
|
|
seen_evidence: set[str] = set()
|
|
gold_doc_id = f"GOLD::{qid}"
|
|
for raw_rank, source in enumerate(sources, 1):
|
|
evidence = str(source.get("evidence") or "")
|
|
normalized = normalize(evidence)
|
|
if not normalized or normalized in seen_evidence:
|
|
continue
|
|
seen_evidence.add(normalized)
|
|
similarity = evidence_similarity(gold_source, evidence)
|
|
doc_id = gold_doc_id if similarity["matched"] else evidence_doc_id(evidence)
|
|
row = {
|
|
"rank": len(ranked) + 1,
|
|
"raw_rank": raw_rank,
|
|
"doc_id": doc_id,
|
|
"paper_id": paper_id(source),
|
|
"evidence": evidence,
|
|
"source": source.get("source", ""),
|
|
"relation": source.get("relation", ""),
|
|
"target": source.get("target", ""),
|
|
**similarity,
|
|
}
|
|
ranked.append(row)
|
|
audit.append(row)
|
|
return ranked, audit
|
|
|
|
|
|
def metric_key(measure: Any) -> str:
|
|
return str(measure)
|
|
|
|
|
|
def write_csv(path: Path, rows: list[dict[str, Any]]) -> None:
|
|
if not rows:
|
|
return
|
|
with path.open("w", encoding="utf-8-sig", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
|
|
writer.writeheader()
|
|
writer.writerows(rows)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--results-dir", required=True)
|
|
parser.add_argument("--questions", required=True)
|
|
parser.add_argument("--judgments", required=True)
|
|
parser.add_argument("--output-dir", required=True)
|
|
args = parser.parse_args()
|
|
|
|
results_dir = Path(args.results_dir)
|
|
questions = json.loads(Path(args.questions).read_text(encoding="utf-8-sig"))
|
|
question_by_id = {str(item["id"]): item for item in questions}
|
|
judgment_rows = list(
|
|
csv.DictReader(open(args.judgments, encoding="utf-8-sig"))
|
|
)
|
|
judgment_by_key = {
|
|
(row["config_id"], row["question_id"]): row for row in judgment_rows
|
|
}
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
(output_dir / "trec_runs").mkdir(exist_ok=True)
|
|
|
|
qrels = [
|
|
Qrel(str(item["id"]), f"GOLD::{item['id']}", 1) for item in questions
|
|
]
|
|
qrels_path = output_dir / "known_item_qrels.tsv"
|
|
with qrels_path.open("w", encoding="utf-8", newline="") as handle:
|
|
for item in questions:
|
|
qid = str(item["id"])
|
|
handle.write(f"{qid}\t0\tGOLD::{qid}\t1\n")
|
|
|
|
measures = [
|
|
P @ 5,
|
|
P @ 10,
|
|
P @ 30,
|
|
Recall @ 5,
|
|
Recall @ 10,
|
|
Recall @ 30,
|
|
RR @ 30,
|
|
nDCG @ 5,
|
|
nDCG @ 10,
|
|
nDCG @ 30,
|
|
]
|
|
summary_rows: list[dict[str, Any]] = []
|
|
query_rows: list[dict[str, Any]] = []
|
|
match_audit: list[dict[str, Any]] = []
|
|
|
|
for config in CONFIGS:
|
|
run: list[ScoredDoc] = []
|
|
diversity: dict[int, list[float]] = {k: [] for k in KS}
|
|
diversity_returned: dict[int, list[float]] = {k: [] for k in KS}
|
|
duplicate_rates: list[float] = []
|
|
retrieved_counts: list[int] = []
|
|
per_question_ranked: dict[str, list[dict[str, Any]]] = {}
|
|
|
|
run_path = output_dir / "trec_runs" / f"{config}.run"
|
|
with run_path.open("w", encoding="utf-8", newline="") as run_handle:
|
|
for item in questions:
|
|
qid = str(item["id"])
|
|
record = json.loads(
|
|
(results_dir / config / f"{qid}.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
response = record.get("response") or {}
|
|
raw_sources = response.get("sources") or []
|
|
ranked, audit = dedupe_ranked_sources(
|
|
qid, str(item.get("source_text") or ""), raw_sources
|
|
)
|
|
per_question_ranked[qid] = ranked
|
|
retrieved_counts.append(len(ranked))
|
|
duplicate_rates.append(
|
|
1 - (len(ranked) / len(raw_sources)) if raw_sources else 0.0
|
|
)
|
|
for row in audit:
|
|
match_audit.append(
|
|
{
|
|
"config_id": config,
|
|
"question_id": qid,
|
|
"gold_source_text": item.get("source_text", ""),
|
|
**row,
|
|
}
|
|
)
|
|
for row in ranked:
|
|
score = 10000.0 - float(row["rank"])
|
|
run.append(ScoredDoc(qid, str(row["doc_id"]), score))
|
|
run_handle.write(
|
|
f"{qid} Q0 {row['doc_id']} {row['rank']} "
|
|
f"{score:.6f} {config}\n"
|
|
)
|
|
for k in KS:
|
|
top = ranked[:k]
|
|
unique_papers = len(
|
|
{str(row["paper_id"]) for row in top if row["paper_id"]}
|
|
)
|
|
diversity[k].append(unique_papers / k)
|
|
diversity_returned[k].append(
|
|
unique_papers / len(top) if top else 0.0
|
|
)
|
|
|
|
aggregate = calc_aggregate(measures, qrels, run)
|
|
summary: dict[str, Any] = {
|
|
"config_id": config,
|
|
"questions": len(questions),
|
|
}
|
|
for measure in measures:
|
|
summary[metric_key(measure)] = float(aggregate[measure])
|
|
for k in KS:
|
|
summary[f"source_diversity@{k}"] = sum(diversity[k]) / len(
|
|
diversity[k]
|
|
)
|
|
summary[f"source_diversity_returned@{k}"] = sum(
|
|
diversity_returned[k]
|
|
) / len(diversity_returned[k])
|
|
summary["mean_exact_duplicate_rate"] = sum(duplicate_rates) / len(
|
|
duplicate_rates
|
|
)
|
|
summary["mean_deduplicated_source_count"] = sum(retrieved_counts) / len(
|
|
retrieved_counts
|
|
)
|
|
|
|
answer_counts = Counter()
|
|
answer_scores: list[float] = []
|
|
for item in questions:
|
|
qid = str(item["id"])
|
|
judgment = judgment_by_key[(config, qid)]
|
|
verdict = judgment["answer_verdict"]
|
|
answer_counts[verdict] += 1
|
|
answer_scores.append(
|
|
{"correct": 1.0, "partial": 0.5}.get(verdict, 0.0)
|
|
)
|
|
ranked = per_question_ranked[qid]
|
|
first_rank = next(
|
|
(
|
|
int(row["rank"])
|
|
for row in ranked
|
|
if row["doc_id"] == f"GOLD::{qid}"
|
|
),
|
|
0,
|
|
)
|
|
query_rows.append(
|
|
{
|
|
"config_id": config,
|
|
"question_id": qid,
|
|
"question": item["question"],
|
|
"answer_verdict": verdict,
|
|
"answer_score": {
|
|
"correct": 1.0,
|
|
"partial": 0.5,
|
|
}.get(verdict, 0.0),
|
|
"evidence_verdict_llm": judgment["evidence_verdict"],
|
|
"known_gold_first_rank": first_rank,
|
|
"known_gold_hit@5": int(bool(first_rank and first_rank <= 5)),
|
|
"known_gold_hit@10": int(
|
|
bool(first_rank and first_rank <= 10)
|
|
),
|
|
"known_gold_hit@30": int(
|
|
bool(first_rank and first_rank <= 30)
|
|
),
|
|
"deduplicated_source_count": len(ranked),
|
|
"unique_papers@10": len(
|
|
{
|
|
str(row["paper_id"])
|
|
for row in ranked[:10]
|
|
if row["paper_id"]
|
|
}
|
|
),
|
|
"judge_reason": judgment["reason"],
|
|
}
|
|
)
|
|
summary["answer_correct"] = answer_counts["correct"]
|
|
summary["answer_partial"] = answer_counts["partial"]
|
|
summary["answer_wrong"] = answer_counts["wrong"]
|
|
summary["answer_insufficient"] = answer_counts["insufficient"]
|
|
summary["answer_correctness_1_0.5_0"] = sum(answer_scores) / len(
|
|
answer_scores
|
|
)
|
|
summary_rows.append(summary)
|
|
|
|
scored_rows: list[dict[str, Any]] = []
|
|
for item in questions:
|
|
qid = str(item["id"])
|
|
judgment = judgment_by_key[(config, qid)]
|
|
verdict = judgment["answer_verdict"]
|
|
record = json.loads(
|
|
(results_dir / config / f"{qid}.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
scored_rows.append(
|
|
{
|
|
"id": qid,
|
|
"question": item["question"],
|
|
"gold_answer": item["gold_answer"],
|
|
"deepseek_answer": (record.get("response") or {}).get(
|
|
"answer", ""
|
|
),
|
|
"score": {"correct": 1.0, "partial": 0.5}.get(
|
|
verdict, 0.0
|
|
),
|
|
"label": verdict,
|
|
"reason": judgment["reason"],
|
|
"model": "DeepSeek V4Flash + RAG",
|
|
"status": "answered",
|
|
}
|
|
)
|
|
config_base = output_dir / f"{config}_same_scale_correctness"
|
|
(config_base.with_suffix(".json")).write_text(
|
|
json.dumps(scored_rows, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
write_csv(config_base.with_suffix(".csv"), scored_rows)
|
|
config_summary = {
|
|
"model": f"DeepSeek V4Flash + RAG ({config})",
|
|
"n": len(scored_rows),
|
|
"total_score": sum(float(row["score"]) for row in scored_rows),
|
|
"answer_correctness": summary["answer_correctness_1_0.5_0"],
|
|
"label_counts": dict(answer_counts),
|
|
"score_counts": dict(
|
|
Counter(str(row["score"]) for row in scored_rows)
|
|
),
|
|
"scoring_note": (
|
|
"Unified 1/0.5/0 scale. Labels come from the calibrated "
|
|
"DeepSeek evidence-aware judge, not hard-coded question IDs."
|
|
),
|
|
}
|
|
(output_dir / f"{config}_same_scale_correctness_summary.json").write_text(
|
|
json.dumps(config_summary, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
write_csv(output_dir / "ir_metrics_summary.csv", summary_rows)
|
|
write_csv(output_dir / "ir_metrics_per_question.csv", query_rows)
|
|
write_csv(output_dir / "gold_evidence_match_audit.csv", match_audit)
|
|
report = {
|
|
"protocol": "known-item retrieval evaluation",
|
|
"implementation": "ir_measures 0.4.1 / pytrec-eval-terrier",
|
|
"qrels": str(qrels_path),
|
|
"ks": KS,
|
|
"matching_rule": {
|
|
"containment": True,
|
|
"sequence_ratio_threshold": 0.82,
|
|
"token_recall_threshold": 0.72,
|
|
"token_precision_threshold": 0.28,
|
|
},
|
|
"important_limitation": (
|
|
"Each question has one known gold source text. Therefore Recall@k "
|
|
"is equivalent to known-item hit rate. Precision@k and nDCG@k "
|
|
"treat unjudged evidence as non-relevant and may underestimate "
|
|
"true retrieval quality. Formal publication-grade P/nDCG requires "
|
|
"pooled relevance judgments over retrieved evidence."
|
|
),
|
|
"summary": summary_rows,
|
|
}
|
|
(output_dir / "ir_evaluation_report.json").write_text(
|
|
json.dumps(report, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|