| |
| 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() |
|
|
