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:
parent
c7158418dd
commit
2d908639e9
|
|
@ -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()
|
||||
Loading…
Reference in New Issue