pipelines/prepare_validation_reference.py

from __future__ import annotations

import argparse
import json
from datetime import datetime
from pathlib import Path
from typing import Any


PROJECT_ROOT = Path(__file__).resolve().parent.parent
RUNS_ROOT = PROJECT_ROOT / "runs" / "validation_reference"
DEFAULT_GOLD_DATASET_PATH = PROJECT_ROOT / "runs" / "gold" / "gold_release_v1" / "gold_dataset.jsonl"


class ValidationReferenceBuilder:
    """从全量 Gold 中提取与 Swift val.jsonl 对齐的参考子集。"""

    def run(
        self,
        gold_dataset_path: Path,
        swift_val_dataset_path: Path,
        output_dir: Path | None = None,
    ) -> dict[str, Any]:
        """执行参考子集构建。"""
        gold_dataset_path = gold_dataset_path.expanduser().resolve()
        swift_val_dataset_path = swift_val_dataset_path.expanduser().resolve()
        run_dir = output_dir.expanduser().resolve() if output_dir else self._build_run_dir()
        run_dir.mkdir(parents=True, exist_ok=True)

        gold_rows = self._load_jsonl(gold_dataset_path)
        val_rows = self._load_jsonl(swift_val_dataset_path)
        val_image_paths = {self._extract_image_path_from_swift_row(row) for row in val_rows}
        subset_rows = [row for row in gold_rows if str(row.get("image_path")) in val_image_paths]
        subset_rows = sorted(subset_rows, key=lambda row: str(row.get("image_path", "")))

        output_path = run_dir / "reference_subset.jsonl"
        summary_path = run_dir / "summary.json"
        report_path = run_dir / "report.md"

        self._write_jsonl(output_path, subset_rows)

        summary = {
            "created_at": datetime.now().isoformat(timespec="seconds"),
            "gold_dataset_path": str(gold_dataset_path),
            "swift_val_dataset_path": str(swift_val_dataset_path),
            "reference_subset_path": str(output_path),
            "summary_path": str(summary_path),
            "report_path": str(report_path),
            "val_image_count": len(val_image_paths),
            "matched_reference_rows": len(subset_rows),
            "missing_reference_rows": max(0, len(val_image_paths) - len(subset_rows)),
        }
        summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
        report_path.write_text(self._build_report(summary), encoding="utf-8")
        return summary

    def _build_run_dir(self) -> Path:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return (RUNS_ROOT / f"validation_reference_{timestamp}").resolve()

    def _load_jsonl(self, path: Path) -> list[dict[str, Any]]:
        rows: list[dict[str, Any]] = []
        with path.open("r", encoding="utf-8") as file:
            for line in file:
                line = line.strip()
                if not line:
                    continue
                rows.append(json.loads(line))
        return rows

    def _write_jsonl(self, path: Path, rows: list[dict[str, Any]]) -> None:
        with path.open("w", encoding="utf-8") as file:
            for row in rows:
                file.write(json.dumps(row, ensure_ascii=False))
                file.write("\n")

    def _extract_image_path_from_swift_row(self, row: dict[str, Any]) -> str:
        metadata = row.get("metadata")
        if isinstance(metadata, dict) and isinstance(metadata.get("image_path"), str):
            return str(metadata["image_path"])
        images = row.get("images")
        if isinstance(images, list) and images and isinstance(images[0], str):
            return str(images[0])
        raise ValueError("Swift val 数据缺少 image_path。")

    def _build_report(self, summary: dict[str, Any]) -> str:
        return "\n".join(
            [
                "# Validation Reference Report",
                "",
                f"- created_at: `{summary['created_at']}`",
                f"- gold_dataset_path: `{summary['gold_dataset_path']}`",
                f"- swift_val_dataset_path: `{summary['swift_val_dataset_path']}`",
                f"- reference_subset_path: `{summary['reference_subset_path']}`",
                f"- val_image_count: `{summary['val_image_count']}`",
                f"- matched_reference_rows: `{summary['matched_reference_rows']}`",
                f"- missing_reference_rows: `{summary['missing_reference_rows']}`",
                "",
            ]
        )


def run_validation_reference_build(
    gold_dataset_path: Path,
    swift_val_dataset_path: Path,
    output_dir: Path | None = None,
) -> dict[str, Any]:
    """运行验证集参考子集构建。"""
    return ValidationReferenceBuilder().run(
        gold_dataset_path=gold_dataset_path,
        swift_val_dataset_path=swift_val_dataset_path,
        output_dir=output_dir,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="从 Gold 中提取与 Swift val.jsonl 对齐的参考子集。")
    parser.add_argument("--gold-dataset-path", default=str(DEFAULT_GOLD_DATASET_PATH), help="全量 Gold 数据集路径")
    parser.add_argument("--swift-val-dataset-path", required=True, help="Swift val.jsonl 路径")
    parser.add_argument("--output-dir", default=None, help="输出目录")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    summary = run_validation_reference_build(
        gold_dataset_path=Path(args.gold_dataset_path),
        swift_val_dataset_path=Path(args.swift_val_dataset_path),
        output_dir=Path(args.output_dir).expanduser().resolve() if args.output_dir else None,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))