#!/usr/bin/env python3 """Convert AnomSeer validation JSONL to Time-RA-style prediction and metric JSON.""" import argparse import json import os from typing import Any, Dict, List from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score from verl.utils.rats_eval import compute_reasoning_metrics from verl.utils.reward_score.rats import NAME_BY_ID TYPE_LABELS = list(NAME_BY_ID) BINARY_LABELS = [0, 1] def load_records(path: str) -> List[Dict[str, Any]]: records = [] with open(path, encoding="utf-8") as input_file: for line_number, line in enumerate(input_file, start=1): line = line.strip() if not line: continue try: records.append(json.loads(line)) except json.JSONDecodeError as exc: raise ValueError(f"Invalid JSON on line {line_number} of {path}") from exc return records def compute_metrics( records: List[Dict[str, Any]], expected_samples: int | None = None, ) -> Dict[str, Any]: valid_pairs = [] missing_ground_truth = 0 for record in records: gt_id = record.get("gt_action_id") pred_id = record.get("pred_action_id") if gt_id is None or int(gt_id) < 0: missing_ground_truth += 1 continue if pred_id is None or int(pred_id) < 0: continue valid_pairs.append((int(pred_id), int(gt_id))) metrics: Dict[str, Any] = { "num_dataset_samples": expected_samples if expected_samples is not None else len(records), "num_prediction_samples": len(records), "num_common_samples": len(records) - missing_ground_truth, "num_valid_samples": len(valid_pairs), "num_invalid_predictions": len(records) - missing_ground_truth - len(valid_pairs), "num_missing_predictions": max((expected_samples or len(records)) - len(records), 0), } if valid_pairs: pred_type = [pred for pred, _ in valid_pairs] gt_type = [gt for _, gt in valid_pairs] pred_binary = [0 if pred == 0 else 1 for pred in pred_type] gt_binary = [0 if gt == 0 else 1 for gt in gt_type] metrics.update({ "type_accuracy": accuracy_score(gt_type, pred_type), "type_precision_macro": precision_score( gt_type, pred_type, labels=TYPE_LABELS, average="macro", zero_division=0, ), "type_recall_macro": recall_score( gt_type, pred_type, labels=TYPE_LABELS, average="macro", zero_division=0, ), "type_f1_macro": f1_score( gt_type, pred_type, labels=TYPE_LABELS, average="macro", zero_division=0, ), "binary_accuracy": accuracy_score(gt_binary, pred_binary), "binary_precision_macro": precision_score( gt_binary, pred_binary, labels=BINARY_LABELS, average="macro", zero_division=0, ), "binary_recall_macro": recall_score( gt_binary, pred_binary, labels=BINARY_LABELS, average="macro", zero_division=0, ), "binary_f1_macro": f1_score( gt_binary, pred_binary, labels=BINARY_LABELS, average="macro", zero_division=0, ), }) else: for key in ( "type_accuracy", "type_precision_macro", "type_recall_macro", "type_f1_macro", "binary_accuracy", "binary_precision_macro", "binary_recall_macro", "binary_f1_macro", ): metrics[key] = 0.0 reasoning_metrics = compute_reasoning_metrics(records) metrics.update(reasoning_metrics) metrics["num_nonempty_thoughts"] = sum( bool(str(record.get("pred_thought", "") or "").strip()) for record in records ) return metrics def build_predictions(records: List[Dict[str, Any]], split: str) -> Dict[str, Any]: predictions = {} for fallback_index, record in enumerate(records): sample_index = record.get("index") sample_key = str(fallback_index if sample_index is None else sample_index) if sample_key in predictions: raise ValueError(f"Duplicate prediction index: {sample_key}") pred_id = record.get("pred_action_id") pred_id = int(pred_id) if pred_id is not None else -1 prediction = { "Thought": str(record.get("pred_thought", "") or ""), "RawResponse": str(record.get("response", "") or ""), "ReferenceThought": str(record.get("gt_thought", "") or ""), "ReferenceActionID": record.get("gt_action_id"), "ReferenceAction": record.get("gt_class"), "Source": record.get("source"), "ImagePath": record.get("image_path"), } if pred_id >= 0: prediction.update({ "ActionID": pred_id, "Label": 0 if pred_id == 0 else 1, "Action": record.get("pred_class") or NAME_BY_ID.get(pred_id), }) else: prediction["ParseError"] = "unrecognized_action" predictions[sample_key] = prediction return {split: predictions} def write_json(path: str, value: Dict[str, Any]) -> None: path = os.path.abspath(os.path.expanduser(path)) os.makedirs(os.path.dirname(path), exist_ok=True) temporary_path = f"{path}.tmp" with open(temporary_path, "w", encoding="utf-8") as output_file: json.dump(value, output_file, ensure_ascii=False, indent=2) output_file.write("\n") os.replace(temporary_path, path) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--input-jsonl", required=True) parser.add_argument("--predictions-output", required=True) parser.add_argument("--metrics-output", required=True) parser.add_argument("--split", default="TSAD_test") parser.add_argument("--checkpoint", default=None) parser.add_argument("--train-file", default=None) parser.add_argument("--test-file", default=None) parser.add_argument("--require-complete", action="store_true") args = parser.parse_args() records = load_records(args.input_jsonl) expected_samples = None if args.test_file and args.test_file.endswith(".parquet"): import pyarrow.parquet as pq expected_samples = pq.ParquetFile(args.test_file).metadata.num_rows if args.require_complete and expected_samples is not None and len(records) != expected_samples: raise ValueError( f"Expected {expected_samples} predictions from {args.test_file}, " f"but found {len(records)} in {args.input_jsonl}" ) predictions = build_predictions(records, args.split) metrics = compute_metrics(records, expected_samples=expected_samples) metrics.update({ "checkpoint": args.checkpoint, "train_file": args.train_file, "test_file": args.test_file, "predictions_file": os.path.abspath(args.predictions_output), }) write_json(args.predictions_output, predictions) write_json(args.metrics_output, metrics) print(f"Saved {len(records)} predictions to {os.path.abspath(args.predictions_output)}") print(f"Saved metrics to {os.path.abspath(args.metrics_output)}") if __name__ == "__main__": main()