File size: 5,712 Bytes
9a069e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/usr/bin/env python
import argparse
import json
import re
from collections import defaultdict
from pathlib import Path
import pyarrow.parquet as pq
from eval_segmentqa_vllm import correct
BUCKETS = {'1': 'or', '2': 'tg', '3': 'agg', '4': 'eu', '5': 'cr'}
TIME_RE = re.compile(r'\d{2}:\d{2}:\d{2}(?:\s*,\s*\d{2}:\d{2}:\d{2})*')
def seconds(value):
hours, minutes, secs = map(int, value.split(':'))
if minutes >= 60 or secs >= 60:
raise ValueError(value)
return hours * 3600 + minutes * 60 + secs
def parse_times(value):
value = re.sub(r'<think>.*?</think>', '', str(value or ''), flags=re.S).strip()
if not TIME_RE.fullmatch(value):
return None
try:
return sorted(seconds(item.strip()) for item in value.split(','))
except ValueError:
return None
def summarize(decisions, total):
by_bucket = {}
for bucket in ('agg', 'tg', 'or', 'eu', 'cr'):
values = decisions[bucket]
passed = sum(values)
by_bucket[bucket] = {'correct': passed, 'total': len(values),
'accuracy': passed / len(values) if values else 0.0}
passed = sum(sum(values) for values in decisions.values())
return {'correct': passed, 'total': total, 'accuracy': passed / total if total else 0.0,
'by_bucket': by_bucket}
def summarize_formats(decisions):
by_format = {}
for fmt, values in decisions.items():
passed = sum(values)
by_format[fmt] = {'correct': passed, 'total': len(values),
'accuracy': passed / len(values) if values else 0.0}
return by_format
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--predictions', required=True)
parser.add_argument('--references', required=True)
parser.add_argument('--parquet', action='append', required=True)
parser.add_argument('--local-output', required=True)
parser.add_argument('--time-output', required=True)
args = parser.parse_args()
predictions = [json.loads(line) for line in Path(args.predictions).read_text().splitlines() if line.strip()]
references = {str(row['id']): row for row in
(json.loads(line) for line in Path(args.references).read_text().splitlines() if line.strip())}
metadata = {}
for path in args.parquet:
metadata.update({str(row['id']): row for row in pq.read_table(path).to_pylist()})
ids = [str(row['id']) for row in predictions]
expected = set(references)
actual = set(ids)
completeness = {
'predictions': len(predictions),
'unique_ids': len(actual),
'missing': len(expected - actual),
'duplicate': len(ids) - len(actual),
'unexpected': len(actual - expected),
'failed': sum(bool(row.get('error') or row.get('failed')) for row in predictions),
'skipped': sum(bool(row.get('skipped')) for row in predictions),
'empty_predictions': sum(not str(row.get('prediction', '')).strip() for row in predictions),
}
if completeness != {'predictions': len(expected), 'unique_ids': len(expected), 'missing': 0,
'duplicate': 0, 'unexpected': 0, 'failed': 0, 'skipped': 0,
'empty_predictions': 0}:
raise RuntimeError(f'incomplete predictions: {completeness}')
if set(metadata) != expected:
raise RuntimeError(f'official parquet/reference ID mismatch: parquet={len(metadata)} refs={len(expected)}')
local, tolerant = defaultdict(list), defaultdict(list)
local_formats, tolerant_formats = defaultdict(list), defaultdict(list)
time_parse_errors = 0
for row in predictions:
qid = str(row['id'])
reference, meta = references[qid], metadata[qid]
bucket = BUCKETS[str(meta['primary_capability'])[0]]
local_result = correct(row['prediction'], reference['answer'], reference['answer_format'])
local[bucket].append(local_result)
local_formats[reference['answer_format']].append(local_result)
tolerant_result = local_result
if reference['answer_format'] == 'time':
pred_times, gold_times = parse_times(row['prediction']), parse_times(reference['answer'])
if pred_times is None:
time_parse_errors += 1
tolerant_result = False
else:
duration = seconds(meta['timestamp_end']) - seconds(meta['timestamp_start'])
tolerance = min(5.0, 1.0 + duration * 4.0 / 360.0)
tolerant_result = len(pred_times) == len(gold_times) and all(
abs(pred - gold) <= tolerance for pred, gold in zip(pred_times, gold_times))
tolerant[bucket].append(tolerant_result)
tolerant_formats[reference['answer_format']].append(tolerant_result)
local_metrics = {'protocol': 'scripts/eval_segmentqa_vllm.py::correct',
**completeness, **summarize(local, len(predictions)),
'by_format': summarize_formats(local_formats)}
tolerant_metrics = {
'protocol': 'time-tolerance-only',
'time_tolerance_seconds': 'min(5, 1 + segment_duration * 4 / 360)',
'non_time_protocol': 'scripts/eval_segmentqa_vllm.py::correct',
'time_parse_errors': time_parse_errors,
**completeness,
**summarize(tolerant, len(predictions)),
'by_format': summarize_formats(tolerant_formats),
}
Path(args.local_output).write_text(json.dumps(local_metrics, indent=2))
Path(args.time_output).write_text(json.dumps(tolerant_metrics, indent=2))
print(json.dumps({'local': local_metrics, 'time_tolerance': tolerant_metrics}, indent=2))
if __name__ == '__main__':
main()
|