Source code for medvision_bm.benchmark.summarize_AD_task

"""Summarize Angle/Distance (AD) benchmark results across models.

Reads the parsed ``*.jsonl`` prediction files under each model's ``parsed``
subdirectory, groups samples by dataset/metric label, and aggregates MAE, MRE,
nMAE, success rate, and threshold-based accuracies. Per-model summaries are
written as JSON (raw values and metrics) and a formatted cross-model report is
saved to a text file.

Run as a CLI; see :func:`parse_args` for the accepted arguments.
"""

import argparse
import ast
import glob
import json
import multiprocessing
import os
import re

import numpy as np

from medvision_bm.medvision_lmms_eval.lmms_eval.tasks.medvision.medvision_utils import (
    _compute_physical_diagonal,
)
from medvision_bm.utils.configs import (
    AD_NEAR_ZERO_GT_THRESHOLD,
    SUMMARY_FILENAME_AD_METRICS,
    SUMMARY_FILENAME_AD_VALUES,
)
from medvision_bm.utils.parse_utils import convert_numpy_to_python, get_subfolders




[docs]
def cal_metrics_AD_task(results):
    """Calculate metrics for a single Angle/Distance (AD) estimation sample.

    Parses the model's predicted scalar from ``filtered_resps`` and compares it
    with the ground-truth ``target`` to compute the mean absolute error (MAE),
    mean relative error (MRE), and a success flag. When distance metadata is
    available, a normalized MAE (nMAE) is either read from a precomputed value
    or reconstructed by dividing the MAE by the physical image diagonal.

    Args:
        results (dict): A single sample with keys:

            - ``filtered_resps`` (list): One-element list holding the
              prediction string.
            - ``target`` (str): Ground-truth value, parsed with
              ``ast.literal_eval``.
            - ``doc_meta`` (dict, optional): Metadata used to derive nMAE,
              including any precomputed ``nmae_precomputed`` entry and the
              ``metric_type``/scale fields.

    Returns:
        dict: Metric entries whose keys match the task YAML ``metric`` fields:

            - ``avgMAE``: ``{"MAE": float, "success": bool}``
            - ``avgMRE``: ``{"MRE": float, "success": bool}``
            - ``SuccessRate``: ``{"success": bool}``
            - ``nMAE``: ``{"NMAE": float, "success": bool}``

    Note:
        MAE/MRE are ``np.nan`` and ``success`` is ``False`` when the prediction
        cannot be parsed or does not contain exactly one value. The metric key
        names must match the ``metric`` field in the task's YAML configuration.
    """
    pred = results["filtered_resps"][0]
    target_metrics = np.array(ast.literal_eval(results["target"]))
    try:
        # Parse prediction: expect comma-separated values, convert to float32
        prd_parts = pred.strip().split(",")
        pred_metrics = np.array([np.float32(part.strip()) for part in prd_parts])

        # Validate: prediction must be a single value (not multiple)
        if len(pred_metrics) != 1:
            mean_absolute_error = np.nan
            mean_relative_error = np.nan
            success = False
        else:
            # Calculate errors
            absolute_error = np.abs(pred_metrics - target_metrics)
            mean_absolute_error = np.mean(absolute_error)
            # Add small epsilon to avoid division by zero
            mean_relative_error = np.mean(absolute_error / (target_metrics + 1e-15))
            success = True
    except:
        # Handle any parsing or computation errors
        mean_absolute_error = np.nan
        mean_relative_error = np.nan
        success = False

    doc_meta = results.get("doc_meta")
    nmae_precomputed = doc_meta.get("nmae_precomputed") if doc_meta else None

    if nmae_precomputed is not None:
        nmae_raw = nmae_precomputed.get("NMAE")
        nmae = float(nmae_raw) if nmae_raw is not None else np.nan
        nmae_success = bool(nmae_precomputed.get("success", False)) and np.isfinite(
            nmae
        )
    elif success and doc_meta is not None and doc_meta.get("metric_type") == "distance":
        # Fallback: recompute diagonal from stored or hash-derived scale.
        # Tier 2 (pixel_size_scale present): uses the scale factor stored at eval time — guaranteed correct.
        # Tier 3 (pixel_size_scale absent, old pre-fix JSONL): hash-based derivation; requires
        #   MEDVISION_SCALED_PS_LOW/HIGH env vars to match eval-time values for scaledPS tasks.
        pixel_size_scale = doc_meta.get("pixel_size_scale")
        try:
            diagonal = _compute_physical_diagonal(
                doc_meta,
                scale_mode=doc_meta.get("scale_mode"),
                explicit_scale=pixel_size_scale,
            )
            nmae = float(mean_absolute_error) / diagonal
            nmae_success = True
        except Exception:
            nmae = np.nan
            nmae_success = False
    else:
        nmae = np.nan
        nmae_success = False

    return {
        "avgMAE": {"MAE": mean_absolute_error, "success": success},
        "avgMRE": {"MRE": mean_relative_error, "success": success},
        "SuccessRate": {"success": success},
        "nMAE": {"NMAE": nmae, "success": nmae_success},
    }




def _initialize_metric_counters_AD_task():
    """
    Initialize all metric counters for AD task aggregation.

    Returns:
        dict: Dictionary containing initialized counters:
            - 'sum_MAE': Cumulative sum of MAE values
            - 'sum_MRE': Cumulative sum of MRE values
            - 'num_success': Count of successful predictions
            - 'count_valid_AE': Count of samples with valid absolute error
            - 'count_valid_RE': Count of samples with valid relative error
            - 'count_AE_thresholds': List of 10 bins for AE distribution
            - 'count_RE_thresholds': List of 10 bins for RE distribution

    Note:
        Threshold bins: [0.0-0.1), [0.1-0.2), ..., [0.8-0.9), [0.9-∞)
    """
    return {
        "sum_MAE": 0,
        "sum_MRE": 0,
        "sum_NMAE": 0,
        "num_success": 0,
        "count_valid_AE": 0,
        "count_valid_RE": 0,
        "count_valid_NMAE": 0,
        "count_AE_thresholds": [0] * 10,
        "count_RE_thresholds": [0] * 10,
    }


def _update_mae_counters(value, counters):
    """
    Update Mean Absolute Error (MAE) related counters.

    Updates cumulative sum, valid count, and threshold distribution bins.

    Args:
        value (float): MAE value to process
        counters (dict): Counter dictionary to update in-place

    Note:
        - Only processes non-NaN values
        - Threshold buckets: i covers [i*0.1, (i+1)*0.1), bucket 9 covers [0.9, ∞)
    """
    # Guard against non-finite values (inf, -inf, nan) that cannot be
    # converted to int or meaningfully summed.
    if np.isfinite(value):
        counters["sum_MAE"] += value
        counters["count_valid_AE"] += 1

        # Determine which threshold bucket this value belongs to
        threshold_index = min(int(value * 10), 9)
        counters["count_AE_thresholds"][threshold_index] += 1


def _update_mre_counters(value, counters):
    """
    Update Mean Relative Error (MRE) related counters.

    Updates cumulative sum, valid count, and threshold distribution bins.

    Args:
        value (float): MRE value to process
        counters (dict): Counter dictionary to update in-place

    Note:
        - Only processes non-NaN values
        - Threshold buckets: i covers [i*0.1, (i+1)*0.1), bucket 9 covers [0.9, ∞)
    """
    # Guard against non-finite values (inf, -inf, nan) that cannot be
    # converted to int or meaningfully summed.
    if np.isfinite(value):
        counters["sum_MRE"] += value
        counters["count_valid_RE"] += 1

        # Determine which threshold bucket this value belongs to
        threshold_index = min(int(value * 10), 9)
        counters["count_RE_thresholds"][threshold_index] += 1


def _update_metric_counters_AD_task(metrics_dict, counters):
    """
    Update all metric counters based on calculated metrics for a single sample.

    This is a convenience function that updates MAE, MRE, and success counters
    in a single call.

    Args:
        metrics_dict (dict): Dictionary containing avgMAE, avgMRE, and SuccessRate metrics
        counters (dict): Counter dictionary to update in-place
    """
    # Update MAE counters (skip nan and inf — non-finite values are unparseable
    # responses and should not contribute to the running average)
    if np.isfinite(metrics_dict["avgMAE"]["MAE"]):
        _update_mae_counters(metrics_dict["avgMAE"]["MAE"], counters)

    # Update MRE counters
    if np.isfinite(metrics_dict["avgMRE"]["MRE"]):
        _update_mre_counters(metrics_dict["avgMRE"]["MRE"], counters)

    # Update nMAE counters
    nmae_entry = metrics_dict.get("nMAE", {})
    if nmae_entry.get("success") and np.isfinite(nmae_entry.get("NMAE", np.nan)):
        counters["sum_NMAE"] += nmae_entry["NMAE"]
        counters["count_valid_NMAE"] += 1

    # Update success count
    counters["num_success"] += metrics_dict["SuccessRate"]["success"]


def _calculate_final_metrics_AD_task(counters, count_total):
    """
    Calculate final aggregated metrics from accumulated counters.

    Computes average errors, success rate, and cumulative threshold-based accuracies.

    Args:
        counters (dict): Dictionary of accumulated counters
        count_total (int): Total number of samples processed

    Returns:
        dict: Dictionary containing computed metrics:
            - 'avgMAE': Average mean absolute error across all valid samples
            - 'avgMRE': Average mean relative error across all valid samples
            - 'SuccessRate': Percentage of successful predictions (0.0 to 1.0)
            - 'MAE<k': Cumulative accuracy for k in [0.1, 0.2, ..., 1.0]
            - 'MRE<k': Cumulative accuracy for k in [0.1, 0.2, ..., 1.0]
            - 'num_samples': Total sample count

    Note:
        - Returns np.nan for averages if no valid samples exist
        - Cumulative accuracies represent proportion of samples below threshold k
    """
    task_metrics = {
        "avgMAE": (
            counters["sum_MAE"] / counters["count_valid_AE"]
            if counters["count_valid_AE"] > 0
            else np.nan
        ),
        "avgMRE": (
            counters["sum_MRE"] / counters["count_valid_RE"]
            if counters["count_valid_RE"] > 0
            else np.nan
        ),
        "avgNMAE": (
            counters["sum_NMAE"] / counters["count_valid_NMAE"]
            if counters["count_valid_NMAE"] > 0
            else np.nan
        ),
        "SuccessRate": (
            counters["num_success"] / count_total if count_total > 0 else 0.0
        ),
        "num_samples": count_total,
    }

    # Calculate cumulative threshold-based accuracies
    keys = ["RE"]
    for key in keys:
        for k in range(1, 11):
            # Sum counts from bucket 0 to bucket k-1 (inclusive)
            cumulative_count = sum(counters[f"count_{key}_thresholds"][0:k])
            task_metrics[f"M{key}<{k / 10:.1f}"] = (
                cumulative_count / count_total if count_total > 0 else 0.0
            )

    return task_metrics




[docs]
def process_label_group(label, data):
    """
    Helper function to process metrics for a single label group.
    Used for both sequential and parallel processing.
    """
    targets = data["targets"]
    responses = data["responses"]
    doc_metas = data.get("doc_metas", [None] * len(targets))

    # Skip empty groups
    if not targets or not responses:
        return label, None

    # Initialize counters for this group
    counters = _initialize_metric_counters_AD_task()
    count_total = 0

    # Process each target-response pair
    for target, response, doc_meta in zip(targets, responses, doc_metas):
        try:
            parsed = json.loads(target)
            gt_val = float(parsed[0]) if isinstance(parsed, list) else float(parsed)
        except Exception:
            gt_val = None
        if gt_val is not None and gt_val < AD_NEAR_ZERO_GT_THRESHOLD:
            continue

        count_total += 1
        mock_results = {
            "filtered_resps": [response],
            "target": target,
            "doc_meta": doc_meta,
        }
        metrics_dict = cal_metrics_AD_task(mock_results)
        _update_metric_counters_AD_task(metrics_dict, counters)

    if count_total == 0:
        return label, None

    # Calculate and store final metrics for this label
    task_metrics = _calculate_final_metrics_AD_task(counters, count_total)

    return label, task_metrics






[docs]
def calculate_summary_metrics_per_anatomy_AD_task(all_data, processes=None):
    """
    Calculate summary metrics grouped by label (anatomy/metric type).

    This function aggregates predictions by label and computes comprehensive metrics
    for each unique label (e.g., dataset_metricType_metricKey combinations).

    Args:
        all_data (list): List of dictionaries, each containing:
            - 'label' (str): Label identifier (e.g., "FeTA24_distance_BPD")
            - 'targets' (str): Ground truth value
            - 'responses' (list): Model predictions
        processes (int, optional): Number of processes to use for parallel calculation.

    Returns:
        dict: Dictionary mapping each label to its computed metrics:
            - avgMAE, avgMRE, SuccessRate
            - MAE<k and MRE<k for k in [0.1, 0.2, ..., 1.0]
            - num_samples

    Note:
        Groups data by label before computing metrics to enable per-anatomy analysis
    """
    # Group data by label
    grouped_data = {}
    for item in all_data:
        label = item["label"]
        target = item["targets"]
        response = item["responses"][0]
        doc_meta = item.get("doc_meta")

        if label not in grouped_data:
            grouped_data[label] = {"targets": [], "responses": [], "doc_metas": []}

        grouped_data[label]["targets"].append(target)
        grouped_data[label]["responses"].append(response)
        grouped_data[label]["doc_metas"].append(doc_meta)
    summary_metrics = {}

    # Prepare items for processing
    items = list(grouped_data.items())

    if processes is not None and processes > 1:
        print(f"Calculating metrics with {processes} processes...")
        with multiprocessing.Pool(processes=processes) as pool:
            results = pool.starmap(process_label_group, items)
    else:
        results = [process_label_group(label, data) for label, data in items]

    # Collect results
    for label, task_metrics in results:
        if task_metrics is not None:
            # Initialize the label entry if it doesn't exist
            if label not in summary_metrics:
                summary_metrics[label] = {}
            summary_metrics[label] = task_metrics

    return summary_metrics






[docs]
def find_and_group_jsonl_files(model_path):
    """Find and group JSONL files in a model directory by dataset and task.

    Different datasets use different grouping strategies:

    - Ceph-Biometrics-400: each file is kept separate (one file per group).
    - FeTA24: all files for the same task are grouped together.

    Args:
        model_path (str): Path to the directory containing JSONL files.

    Returns:
        dict: Dictionary mapping group keys to lists of file paths. A
            Ceph-Biometrics-400 file maps to a single-element list keyed by its
            filename, while all FeTA24 files for a task map to the combined key
            ``"FeTA24_BiometricsFromLandmarks_Task01_combined"``.
    """
    # Find all JSONL files in the model folder (exclude analysis output files)
    jsonl_files = [
        f
        for f in glob.glob(os.path.join(model_path, "*.jsonl"))
        if not ("_proc_acc" in os.path.basename(f) or "_eq_acc" in os.path.basename(f))
    ]

    # Group files by dataset and task
    grouped_files = {}

    for jsonl_file in jsonl_files:
        filename = os.path.basename(jsonl_file)

        # Handle Ceph-Biometrics-400: no grouping, one file per key
        if "Ceph-Biometrics-400" in filename:
            grouped_files[filename] = [jsonl_file]

        # Handle FeTA24: group all files for the same task together
        elif "FeTA24_BiometricsFromLandmarks_Task01" in filename:
            key = "FeTA24_BiometricsFromLandmarks_Task01_combined"
            if key not in grouped_files:
                grouped_files[key] = []
            grouped_files[key].append(jsonl_file)

    return grouped_files






[docs]
def process_jsonl_file(
    jsonl_path,
    limit,
):
    """
    Process a single JSONL file and extract label, target, and response data.

    Each line in the JSONL file represents one prediction sample. This function
    extracts the biometric profile information and constructs a structured label.

    Args:
        jsonl_path (str): Path to the JSONL file
        limit (int, optional): Maximum number of samples to process. None for all samples.

    Returns:
        list: List of dictionaries, each containing:
            - 'label' (str): Constructed label (dataset_metricType_metricKey)
            - 'targets' (str): Ground truth value
            - 'responses' (list): Model predictions

    Raises:
        ValueError: If JSON parsing fails
        AssertionError: If metric_key is None
    """
    results = []
    # Extract dataset name from filename pattern 'samples_{dataset_name}_'
    match = re.search(r"samples_([^_]+)_", os.path.basename(jsonl_path))
    dataset_name = match.group(1)

    count = 0
    with open(jsonl_path, "r") as f:
        for line_idx, line in enumerate(f):
            # Skip empty lines
            if not line.strip():
                continue

            try:
                data = json.loads(line.strip())
                # Skip empty data
                if not data:
                    continue

                doc = data.get("doc", {})

                filtered_resps = data.get("filtered_resps")
                target = data.get("target")

                # Extract biometric profile information
                biometric_profile = doc.get("biometric_profile", {})
                metric_type = biometric_profile.get("metric_type", "")
                metric_key = biometric_profile.get("metric_key")
                assert (
                    metric_key is not None
                ), f"metric_key is None in line {line_idx + 1} of {jsonl_path}"

                # Construct label: dataset_metricType_metricKey
                # Example: "FeTA24_distance_BPD" or "Ceph-Biometrics-400_angle_SNA"
                label_name = f"{dataset_name}_{metric_type}_{metric_key}"
                scale_mode = (
                    "anisotropic"
                    if "scaledPS" in os.path.basename(jsonl_path)
                    else None
                )
                doc_meta = {
                    "image_file": doc.get("image_file"),
                    "slice_dim": doc.get("slice_dim"),
                    "slice_idx": doc.get("slice_idx"),
                    "image_size_2d": doc.get("image_size_2d"),
                    "metric_type": metric_type,
                    "scale_mode": scale_mode,
                    "nmae_precomputed": data.get("nMAE"),
                    "taskID": doc.get("taskID"),
                    "label": doc.get("label"),
                    "pixel_size_scale": data.get("pixel_size_scale"),
                }
                results.append(
                    {
                        "label": label_name,
                        "targets": target,
                        "responses": filtered_resps,
                        "doc_meta": doc_meta,
                    }
                )

                count += 1
                if limit is not None and count >= limit:
                    break

            except json.JSONDecodeError:
                raise ValueError(f"Error in parsing the JSON file {jsonl_path}")

    return results






[docs]
def process_combined_jsonl_files(jsonl_paths, limit):
    """
    Process multiple JSONL files and combine their data into a single list.

    Useful for combining data from multiple files that belong to the same
    dataset or task group (e.g., multiple FeTA24 files).

    Args:
        jsonl_paths (list): List of paths to JSONL files to combine
        limit (int, optional): Maximum number of samples to process per file.
                              None for all samples.

    Returns:
        list: Combined list of dictionaries with 'label', 'targets', and 'responses' keys
    """
    combined_data = []

    for jsonl_path in jsonl_paths:
        file_data = process_jsonl_file(jsonl_path, limit)
        combined_data.extend(file_data)

    return combined_data






[docs]
def process_parsed_file_in_model_folder(
    model_dir,
    limit=None,
    processes=None,
):
    """Process all parsed JSONL files in a model folder and write AD summaries.

    Steps performed:

    1. Find and group JSONL files in the ``parsed`` subdirectory.
    2. Extract targets and model predictions from each file.
    3. Calculate comprehensive metrics per anatomy/metric type.
    4. Save two JSON files in the ``parsed`` directory: a values file with the
       raw targets and predictions, and a metrics file with the aggregated
       metrics per label.

    If the ``parsed`` subdirectory does not exist, a message is printed and the
    function returns without writing any output.

    Args:
        model_dir (str): Path to the model folder (expected to contain a
            ``parsed`` subdirectory).
        limit (int, optional): Maximum number of samples to process per file.
            None for all samples.
        processes (int, optional): Number of processes to use for parallel
            calculation.
    """
    # Locate parsed files directory
    parsed_files_dir = os.path.join(model_dir, "parsed")

    # # Option 1: Early exit if parsed directory does not exist
    # assert os.path.exists(
    #     parsed_files_dir
    # ), f"Parsed files directory does not exist: {parsed_files_dir}"

    # Option 2: Warning and skip if parsed directory does not exist
    if not os.path.exists(parsed_files_dir):
        print(f"Parsed files directory does not exist: {parsed_files_dir}, skipping...")
        return

    grouped_files = find_and_group_jsonl_files(parsed_files_dir)

    # Collect all data from the parsed JSONL files
    all_data = []
    for group_name, file_paths in grouped_files.items():
        if len(file_paths) == 1:
            # Single file processing (e.g., Ceph-Biometrics-400)
            file_data = process_jsonl_file(file_paths[0], limit)
        else:
            # Combined file processing (e.g., multiple FeTA24 files)
            file_data = process_combined_jsonl_files(file_paths, limit)

        all_data.extend(file_data)

    # Skip if no valid data was collected
    if not all_data:
        print(f"No valid data found in {parsed_files_dir}, skipping...")
        return

    # Calculate summary metrics per anatomy
    summary_metrics = calculate_summary_metrics_per_anatomy_AD_task(
        all_data, processes=processes
    )

    # Save raw values JSON file (targets and predictions for each sample)
    values_filename = (
        SUMMARY_FILENAME_AD_VALUES
        if limit is None
        else f"{SUMMARY_FILENAME_AD_VALUES.removesuffix('.json')}_limit{limit}.json"
    )
    values_path = os.path.join(parsed_files_dir, values_filename)
    with open(values_path, "w") as f:
        json.dump(convert_numpy_to_python(all_data), f, indent=2)
    print(f"Saved target and model-predicted values to {values_path}")

    # Save aggregated metrics JSON file (metrics per label)
    metrics_filename = (
        SUMMARY_FILENAME_AD_METRICS
        if limit is None
        else f"{SUMMARY_FILENAME_AD_METRICS.removesuffix('.json')}_limit{limit}.json"
    )
    metrics_path = os.path.join(parsed_files_dir, metrics_filename)
    with open(metrics_path, "w") as f:
        json.dump(convert_numpy_to_python(summary_metrics), f, indent=2)
    print(f"Saved summary metrics to {metrics_path}")








def _process_task_directory(
    task_dir, limit, processes=None, skip_model_wo_parsed_files=False
):
    """
    Process all model directories within a task directory.

    This function orchestrates the entire processing pipeline:
    1. Finds all model subdirectories in the task directory
    2. Processes each model's JSONL files
    3. Generates summary metrics across all models

    Args:
        task_dir (str): Path to the task directory containing model folders
        limit (int, optional): Maximum number of samples to process per file
        processes (int, optional): Number of processes to use for parallel calculation
        skip_model_wo_parsed_files (bool): Whether to skip model directories
                                          without parsed folders
    """
    # Get list of model folders within task_dir
    model_dirs = get_subfolders(task_dir)

    # Loop over each model directory and process JSONL files
    for model_dir in model_dirs:
        # Skip if parsed folder doesn't exist and flag is set
        parsed_files_dir = os.path.join(model_dir, "parsed")
        if skip_model_wo_parsed_files and not os.path.exists(parsed_files_dir):
            print(f"\nSkipping model directory (no parsed folder): {model_dir}")
            continue

        print(f"\nProcessing model directory: {model_dir}")
        process_parsed_file_in_model_folder(model_dir, limit, processes=processes)

    # Print summary metrics at the end
    print_model_summaries(task_dir, limit, skip_model_wo_parsed_files)


def _process_single_model_directory(model_dir, limit, processes=None):
    """
    Process a single model directory.

    Convenience function for processing just one model without generating
    cross-model summaries.

    Args:
        model_dir (str): Path to the model directory
        limit (int, optional): Maximum number of samples to process per file
        processes (int, optional): Number of processes to use for parallel calculation
    """
    print(f"\nProcessing model directory: {model_dir}")
    process_parsed_file_in_model_folder(model_dir, limit, processes=processes)




[docs]
def main(**kwargs):
    """Process AD model folders based on the provided arguments.

    Supports two modes:

    1. **task_dir mode**: process all model directories within a task
       directory and generate a cross-model summary.
    2. **model_dir mode**: process a single model directory in isolation.

    Args:
        task_dir (str, optional): Path to task directory (mutually exclusive
            with model_dir).
        model_dir (str, optional): Path to model directory (mutually exclusive
            with task_dir).
        limit (int, optional): Maximum number of samples to process per file.
            None for all samples.
        skip_model_wo_parsed_files (bool): Whether to skip model directories
            without parsed folders (task_dir mode only).
        processes (int, optional): Number of processes to use for parallel
            calculation.

    Raises:
        ValueError: If neither task_dir nor model_dir is provided.
    """
    task_dir = kwargs.get("task_dir")
    model_dir = kwargs.get("model_dir")
    limit = kwargs.get("limit")
    skip_model_wo_parsed_files = kwargs.get("skip_model_wo_parsed_files", False)
    processes = kwargs.get("processes")

    if task_dir is not None:
        print(
            f"Using task_dir: {task_dir}\nModel directories within this folder will be looped over."
        )
        _process_task_directory(
            task_dir,
            limit,
            processes=processes,
            skip_model_wo_parsed_files=skip_model_wo_parsed_files,
        )

    elif model_dir is not None:
        print(
            f"Using model_dir: {model_dir}\nProcessing all JSONL files within this directory."
        )
        _process_single_model_directory(model_dir, limit, processes=processes)

    else:
        raise ValueError("Either 'task_dir' or 'model_dir' must be provided.")






[docs]
def parse_args():
    """
    Parse and validate command line arguments.

    Returns:
        argparse.Namespace: Parsed arguments with the following attributes:
            - task_dir: Path to task directory (or None)
            - model_dir: Path to model directory (or None)
            - limit: Sample limit (or None for all samples)
            - skip_model_wo_parsed_files: Boolean flag
            - processes: Number of worker processes (or None)

    Raises:
        SystemExit: If arguments are invalid (via parser.error)
    """
    parser = argparse.ArgumentParser(
        description="Process model folders and generate summary metrics."
    )
    parser.add_argument(
        "--task_dir",
        type=str,
        help="Path to the task directory containing model result folders.",
    )
    parser.add_argument(
        "--model_dir",
        type=str,
        help="Path to a specific model directory containing JSONL files.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Limit the number of samples to process per JSONL file. If not set, processes all samples.",
    )
    parser.add_argument(
        "--skip_model_wo_parsed_files",
        action="store_true",
        help="Skip model directories that don't have a 'parsed' folder. Only valid with --task_dir.",
    )
    parser.add_argument(
        "--processes",
        "-p",
        type=int,
        default=None,
        help="Number of worker processes for metric calculation.",
    )

    args = parser.parse_args()

    # Validate that at least one of task_dir or model_dir is provided
    if args.task_dir is None and args.model_dir is None:
        parser.error("Either --task_dir or --model_dir must be provided.")

    # Validate that skip_model_wo_parsed_files is only used with task_dir
    if args.skip_model_wo_parsed_files and args.task_dir is None:
        parser.error("--skip_model_wo_parsed_files can only be used with --task_dir")

    return args




if __name__ == "__main__":
    args_dict = vars(parse_args())
    main(**args_dict)