feat : add ROUGE evaluation script (manual impl, no deps)

- ROUGE-1/2 via n-gram overlap (Counter)
- ROUGE-L via LCS (DP)
- CLI: python scripts/eval/evaluate_rouge.py --data_path ... --output ...
- Library: compute_rouge(ref, cand) -> dict of precision/recall/f1
This commit is contained in:
ViperEkura 2026-07-05 01:15:01 +08:00
parent c7158418dd
commit 2d908639e9
1 changed files with 153 additions and 0 deletions

View File

@ -0,0 +1,153 @@
"""ROUGE evaluation (manual implementation, no external deps).
Computes ROUGE-1, ROUGE-2, ROUGE-L precision, recall, and F1.
Usage::
# Batch evaluation from JSONL (each line: {"reference": ..., "candidate": ...})
python scripts/eval/evaluate_rouge.py --data_path preds.jsonl --output results.json
# As a library
from scripts.eval.evaluate_rouge import compute_rouge
scores = compute_rouge("the cat sat on the mat", "the cat sat")
"""
import argparse
import json
from collections import Counter
from typing import Dict, List, Tuple
def _tokenize(text: str) -> List[str]:
return text.split()
def _ngrams(tokens: List[str], n: int) -> Counter:
return Counter(zip(*[tokens[i:] for i in range(n)]))
def _lcs(x: List[str], y: List[str]) -> int:
m, n = len(x), len(y)
dp = [[0] * (n + 1) for _ in range(m + 1)]
for i in range(1, m + 1):
xi = x[i - 1]
dpi = dp[i]
dpi_1 = dp[i - 1]
for j in range(1, n + 1):
if xi == y[j - 1]:
dpi[j] = dpi_1[j - 1] + 1
else:
dpi[j] = dpi_1[j] if dpi_1[j] > dpi[j - 1] else dpi[j - 1]
return dp[m][n]
def _f1(precision: float, recall: float) -> float:
if precision + recall == 0:
return 0.0
return 2 * precision * recall / (precision + recall)
def _rouge_n(ref_tokens: List[str], cand_tokens: List[str], n: int) -> Dict[str, float]:
ref_ngrams = _ngrams(ref_tokens, n)
cand_ngrams = _ngrams(cand_tokens, n)
overlap = sum((cand_ngrams & ref_ngrams).values())
cand_total = sum(cand_ngrams.values())
ref_total = sum(ref_ngrams.values())
precision = overlap / cand_total if cand_total > 0 else 0.0
recall = overlap / ref_total if ref_total > 0 else 0.0
f1 = _f1(precision, recall)
return {"precision": precision, "recall": recall, "f1": f1}
def _rouge_l(ref_tokens: List[str], cand_tokens: List[str]) -> Dict[str, float]:
lcs_len = _lcs(ref_tokens, cand_tokens)
ref_len = len(ref_tokens)
cand_len = len(cand_tokens)
recall = lcs_len / ref_len if ref_len > 0 else 0.0
precision = lcs_len / cand_len if cand_len > 0 else 0.0
f1 = _f1(precision, recall)
return {"precision": precision, "recall": recall, "f1": f1}
def compute_rouge(
reference: str, candidate: str, n: int = 2
) -> Dict[str, Dict[str, float]]:
"""Compute ROUGE-N (1..n) and ROUGE-L scores.
Returns::
{
"rouge-1": {"precision": ..., "recall": ..., "f1": ...},
"rouge-2": {"precision": ..., "recall": ..., "f1": ...},
"rouge-l": {"precision": ..., "recall": ..., "f1": ...},
}
"""
ref_tokens = _tokenize(reference)
cand_tokens = _tokenize(candidate)
results = {}
for i in range(1, n + 1):
results[f"rouge-{i}"] = _rouge_n(ref_tokens, cand_tokens, i)
results["rouge-l"] = _rouge_l(ref_tokens, cand_tokens)
return results
def evaluate_file(data_path: str) -> Dict:
with open(data_path, "r", encoding="utf-8") as f:
pairs = [json.loads(line) for line in f if line.strip()]
agg = {
k: {"precision": 0.0, "recall": 0.0, "f1": 0.0}
for k in ("rouge-1", "rouge-2", "rouge-l")
}
per_item = []
for item in pairs:
ref = item["reference"]
cand = item["candidate"]
scores = compute_rouge(ref, cand)
per_item.append({**item, "scores": scores})
for k, v in scores.items():
agg[k]["precision"] += v["precision"]
agg[k]["recall"] += v["recall"]
agg[k]["f1"] += v["f1"]
n = len(pairs)
for k in agg:
agg[k] = {m: v / n for m, v in agg[k].items()}
return {"num_samples": n, "aggregate": agg, "per_item": per_item}
def main():
parser = argparse.ArgumentParser(description="ROUGE evaluation")
parser.add_argument(
"--data_path", required=True, help="JSONL with reference/candidate per line"
)
parser.add_argument("--output", type=str, default=None, help="Output JSON path")
args = parser.parse_args()
results = evaluate_file(args.data_path)
agg = results["aggregate"]
print(f"Samples: {results['num_samples']}")
print()
for metric in ("rouge-1", "rouge-2", "rouge-l"):
s = agg[metric]
print(
f" {metric:8s} P={s['precision']:.4f} R={s['recall']:.4f} F1={s['f1']:.4f}"
)
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(results, f, indent=2, ensure_ascii=False)
print(f"\nSaved to {args.output}")
if __name__ == "__main__":
main()