evaluate_data.py

# Description: This script is used to evaluate the results of the model.
# 1) The script loads the data from the specified path.
# 2) It removes unnecessary columns from the data.
# 3) It extracts the split from the data path.
# 4) It iterates over the data and creates the data levels.
# 5) It loads the silver reasoning and interpretability results.
# 5) It calculates the metrics for the split and the tasks.
# 6) It plots the metrics and interpretability heat.
# 7) It saves the results to the specified path.
# 8) It prints the metrics table.
# 9) It categorizes the results into different cases.
# 10) It saves the categorized results to the specified path.

from __future__ import annotations

import argparse
import json
import re
import warnings
from collections import defaultdict
from pathlib import Path

import numpy as np
import pandas as pd

from data.DataLoader import DataLoader
from data.DataProcessor import DataProcessor
from data.DataSaver import DataSaver
from data.utils import format_metrics
from evaluation.utils import extract_split
from inference.DataLevels import Results, Sample, SamplePart, Split, Task, print_metrics
from inference.utils import print_metrics_table
from interpretability.DistractorAttention import (
    DistractorAttentionStats,
    collect_distractor_attention_record,
)
from plots.Plotter import Plotter

PREFIX = Path.cwd()
while PREFIX.name != "research-project":
    PREFIX = PREFIX.parent

supported_single_system_settings = {
    "basic-baseline",
    "baseline",
    "skyline",
}

supported_multi_system_settings = {
    "feedback",
    "speculative_decoding",
    "sd",
}

supported_settings = supported_single_system_settings.union(
    supported_multi_system_settings
)


def remove_unnecessary_columns(
    row: dict[str, str | int | float], headers: dict[str, list[str]]
) -> None:
    """
    Remove unnecessary columns from the row.

    :param row: the row to remove the columns from
    :param headers: the headers of the columns
    :return: None
    """
    unnecessary_columns = [
        "correct",
        "correct?",
        "exact_match_accuracy",
        "soft_match_accuracy",
        "there",
        "verbs",
        "pronouns",
        "not_mentioned",
    ]
    for col in unnecessary_columns:
        if col in row:
            del row[col]
    if not row["silver_reasoning"]:
        del row["silver_reasoning"]
        if "silver_reasoning" in headers["general"]:
            del headers["general"][headers["general"].index("silver_reasoning")]


def extract_split(path) -> str:
    """
    Extract the split from the data path. If the split is not found, return "split".

    :param path: the path to the data
    :return: the split
    """
    for split in ["valid", "test", "train"]:
        if split in path:
            return split
    return "split"


def structure_result(headers_results: list[str], row: dict, version) -> dict[str, list]:
    """
    Structure the result into a dictionary.

    :param headers_results: the headers and results
    :param row: the row of data to structure
    :param version: the version to structure for
    :return: the structured result
    """
    h_patt = re.compile(r"(.+)_(?:after|before)")
    result = [
        (
            (h_patt.match(header)[1], str(row[f"{header}_{version}"]))
            if h_patt.match(header)
            else (header, row[f"{header}_{version}"])
        )
        for header in headers_results
    ]
    return dict(result)


def get_result(
    results_data: list[dict], task_id: int, sample_id: int, part_id: int
) -> dict[str, str] | None:
    """
    Get the result for the task_id, sample_id and part_id.

    :param results_data: the list of result dicts to search
    :param task_id: the task id to match
    :param sample_id: the sample id to match
    :param part_id: the part id to match
    :return: the matching result dict, or None
    """
    for row in results_data:
        if (
            row["task_id"] == task_id
            and row["sample_id"] == sample_id
            and row["part_id"] == part_id
        ):
            return row
    return None


def validate_inputs(run_fn):
    """
    Validate the inputs for the evaluation pipeline.

    :param run_fn: the run function to wrap
    :return: the wrapped function with input validation
    """

    def validation_wrapper(**kwargs):
        experiment = kwargs.get("experiment", "").lower()
        supported_experiments = ["reasoning", "direct_answer"]
        if not experiment:
            raise ValueError(
                f"Please provide an experiment to evaluate from {supported_experiments}"
            )
        if experiment not in supported_experiments:
            raise ValueError(
                f"Experiment '{experiment}' is not supported. "
                f"Please choose either of {supported_experiments}"
            )
        setting = kwargs.get("setting", "").lower()

        if not setting:
            raise ValueError(
                f"Please provide an experiment setting from {supported_settings}"
            )
        if setting not in supported_settings:
            raise ValueError(
                f"Setting not recognized, expected one of: {supported_settings}"
            )
        if not kwargs.get("results_path", ""):
            raise ValueError("Please provide a path to the data for evaluation.")

        reasoning_source = kwargs.get("reasoning_source")
        supported_sources = ["claude", "llama"]
        if reasoning_source is not None and reasoning_source not in supported_sources:
            raise ValueError(
                f"reasoning_source {reasoning_source!r} is not supported. "
                f"Choose one of {supported_sources} or omit it to use the default path."
            )

        filtering_conditions = kwargs.get("filtering_conditions", {})
        if filtering_conditions:
            for attr in filtering_conditions.keys():
                assert hasattr(
                    SamplePart, attr
                ), f"SamplePart does not have the attribute specified in the filtering condition: {attr}"
        return run_fn(**kwargs)

    return validation_wrapper


def _row_normalise_attn_scores(attn_scores: np.ndarray) -> np.ndarray:
    """
    Return a row-normalised copy of a 2-D attention score matrix.

    :param attn_scores: 2-D array of shape ``(output_tokens, num_sentences)``,
                        column-normalised as produced by
                        ``Interpretability.get_attention_scores``.
    :return: row-normalised copy of the same shape; the original is not mutated.
    """
    row_sums = attn_scores.sum(axis=1, keepdims=True)
    safe_row_sums = np.where(row_sums == 0, 1.0, row_sums)
    return attn_scores / safe_row_sums


def _collect_record_with_row_normalised_attn(part, answer_correct, version):
    """
    Temporarily replace the column-normalised ``attn_scores`` on the relevant
    ``InterpretabilityResult`` with a row-normalised copy, call
    ``collect_distractor_attention_record``, then restore the original scores.

    This keeps the column-normalised scores intact for heatmap consumers while
    letting the distractor analysis see the correct normalisation.

    :param part: the ``SamplePart`` passed through to
                 ``collect_distractor_attention_record``.
    :param answer_correct: correctness flag for this part/version.
    :param version: ``"before"`` or ``"after"``.
    :return: the ``DistractorAttentionRecord`` returned by
             ``collect_distractor_attention_record``, or ``None``.
    """
    result_for_version = next((r for r in part.results if r.version == version), None)
    if result_for_version is None:
        return collect_distractor_attention_record(part, answer_correct, version)

    original_scores = getattr(result_for_version.interpretability, "attn_scores", None)
    needs_correction = original_scores is not None and original_scores.ndim == 2

    if needs_correction:
        result_for_version.interpretability.attn_scores = _row_normalise_attn_scores(
            original_scores
        )
    try:
        return collect_distractor_attention_record(part, answer_correct, version)
    finally:
        if needs_correction:
            result_for_version.interpretability.attn_scores = original_scores


@validate_inputs
def run(
    results_path: str,
    save_path: str,
    samples_per_task: int,
    experiment: str,
    setting: str = "baseline",
    filtering_conditions: dict = None,
    create_heatmaps: bool = True,
    verbose: bool = False,
    max_tokens: int | None = None,
    reasoning_source: str | None = None,
) -> None:
    """
    Run the evaluation pipeline.

    :param results_path: path to the data for evaluation
    :param save_path: path to save the results
    :param samples_per_task: number of samples per task the results were run with
    :param experiment: the experiment to evaluate (e.g., "reasoning", "direct_answer")
    :param setting: the setting of the experiment (e.g., "baseline", "feedback")
    :param filtering_conditions: a dictionary of conditions (SamplePart attributes) and values;
                                 all the parts having an attribute with such value will be preserved;
                                 if you need a function result, create a new attribute in
                                 SamplePart __init__ and use it as a filtering condition
    :param create_heatmaps: whether to create heatmaps for the interpretability results
    :param verbose: whether to print the results to the console
    :param max_tokens: the model's generation budget. Passed to
                       :func:`add_completeness_column` so the truncation
                       check can flag rows whose reasoning hit the limit.
                       If ``None`` the truncation check is skipped.
    :param reasoning_source: which silver-reasoning corpus to use as the
                             reference for reasoning quality metrics.
                             ``None`` (default) uses the legacy flat directory
                             ``data/silver_reasoning/``.
                             ``"claude"`` reads from
                             ``data/silver_reasoning/claude/``;
                             ``"llama"`` reads from
                             ``data/silver_reasoning/llama/``.
    :return: None
    """
    print("You are running the evaluation pipeline.", end="\n\n")

    print("Loading data...", end="\n\n")
    if filtering_conditions:
        print("Employing the following filtering conditions for the evaluation:")
        for attr, value in filtering_conditions.items():
            print(f"- {attr} = {value}")
    if reasoning_source:
        print(f"Using silver reasoning source: {reasoning_source!r}", end="\n\n")
    else:
        print("Using default silver reasoning source (legacy path).", end="\n\n")
    loader = DataLoader(
        prefix=PREFIX,
        samples_per_task=samples_per_task,
        filtering_conditions=filtering_conditions,
        reasoning_source=reasoning_source,
    )

    if setting in supported_single_system_settings:
        multi_system = False
    else:
        multi_system = True

    # loaded results in parts with original data, tokens-ids, and interpretability results
    results_data, multi_system = loader.load_results(
        results_path=results_path,
        data_path="../tasks_1-20_v1-2/en-valid/",
        split=extract_split(results_path),
        as_parts=True,
        multi_system=multi_system,
    )

    # maybe loaded_baseline_results is not needed for evaluation
    saver = DataSaver(
        save_to=str(Path(save_path) / "eval"),
        loaded_baseline_results=True if multi_system else False,
    )
    results_file_name = f"{Path(results_path).stem}_upd.csv"
    plotter = Plotter(results_path=saver.run_path, color_map="tab20")
    if filtering_conditions:
        conditions_add = [f"{a}={v}" for a, v in filtering_conditions.items()]
    else:
        conditions_add = []

    print(f"\nLoaded results data for {len(results_data)} tasks.")
    print(f"Loaded {loader.number_of_parts} sample parts created from raw data.")

    distractor_stats: dict[str, DistractorAttentionStats] = defaultdict(
        DistractorAttentionStats
    )
    distractor_stats_per_task: dict[int, dict[str, DistractorAttentionStats]] = (
        defaultdict(lambda: defaultdict(DistractorAttentionStats))
    )

    processor = DataProcessor()

    data_split = extract_split(results_path)
    sample, task, split = None, None, Split(name=data_split, multi_system=multi_system)
    for task_id, samples in results_data.items():
        assert type(task_id) is int
        task = Task(task_id, multi_system=multi_system)

        for sample_id, parts in list(samples.items())[:samples_per_task]:
            assert type(sample_id) is int
            sample = Sample(
                task_id=task_id,
                sample_id=sample_id,
                multi_system=multi_system,
            )

            # Used to store the correct answers for each sample for later evaluation
            for part in parts:
                for version, result in zip(part.versions, part.results):
                    # TODO: add reasoning judgment to part results for it to be saved in the results table
                    if create_heatmaps and not result.interpretability.empty():
                        plotter.draw_heat(
                            result.interpretability,
                            x_label="Sentence Indices",
                            task_id=part.task_id,
                            sample_id=part.sample_id,
                            part_id=part.part_id,
                            version=version,
                            title=f"Attention Map for Task {part.task_id} Sample {part.sample_id} "
                            f"Part {part.part_id} (version: {version}, case: {result.category}, "
                            f"{', '.join(conditions_add)})",
                        )

                sample.add_part(part)
                result = part.get_result()
                # necessary only if we want to addition more columns to our original results
                # otherwise we can just create separate tables or files

                # Distractor marking only needs part.raw and part.supporting_sent_inx,
                # so it runs independently of whether versions and results are paired.
                processor.mark_distractors(part)

                if len(part.versions) != len(part.results):
                    print(
                        f"[WARNING] Skipping malformed part | "
                        f"task={part.task_id} sample={part.sample_id} part={part.part_id}\n"
                        f"versions={part.versions}\n"
                        f"n_results={len(part.results)}"
                    )
                    continue

                distractor_fields: dict = {}
                for version_result in part.results:
                    version = version_result.version

                    if version_result.interpretability.empty():
                        continue

                    result_dict = version_result.get_result()

                    answer_correct = result_dict.get(f"answer_correct_{version}")

                    if answer_correct is None or pd.isna(answer_correct):
                        continue

                    record = _collect_record_with_row_normalised_attn(
                        part=part,
                        answer_correct=answer_correct,
                        version=version,
                    )

                    if record is not None:
                        distractor_stats[version].add(record)
                        distractor_stats_per_task[task_id][version].add(record)
                        # Stash distractor fields to enrich the results row below.
                        distractor_fields[f"attn_distractor_{version}"] = (
                            record.attn_distractor
                        )
                        distractor_fields[f"attn_supporting_{version}"] = (
                            record.attn_supporting
                        )
                        distractor_fields[f"attn_neutral_{version}"] = (
                            record.attn_neutral
                        )
                        if "n_distractors" not in distractor_fields:
                            distractor_fields["n_distractors"] = record.n_distractors
                    else:
                        warnings.warn(
                            f"Empty record collected for task {task_id} sample {sample_id} part {part.part_id} version '{version}'. Check that the part has interpretability data and distractors set."
                        )

                # Enrich the results row with distractor-attention attributes so
                # downstream analyses (e.g. toxic-CoT filtering) have them inline.
                result.update(distractor_fields)
                saver.save_output(
                    data=[result],
                    headers=list(result.keys()),
                    file_name=results_file_name,
                )
            sample.calculate_metrics()
            task.add_sample(sample)

            if verbose:
                sample.print_sample_predictions()
                print_metrics(sample)
            for evaluator, version in zip(sample.evaluators, sample.versions):
                metrics = list(
                    format_metrics(evaluator.get_metrics(as_lists=True)).values()
                )
                print(f"Metrics for {evaluator.level} {version}:", metrics, end="\n\n")

        task.set_results()

        split.add_task(task)
        print(
            f"Added task {task_id} with {len(sample.parts)} parts to split {data_split}."
        )

        task_corr_matrices = task.calculate_metrics()
        if verbose:
            print_metrics(task)

        for version, evaluator, corr_matrix in zip(
            task.versions, task.evaluators, task_corr_matrices.values()
        ):
            # Plot Attention vs Seen Context Lengths for the Task
            plotter.plot_correlation(
                x_data={"Seen context lengths": task.seen_context_lengths},
                y_data=evaluator.parts_attn_on_target.all,
                x_label="Seen Context Lengths",
                y_label="Attention on Target Tokens",
                file_name=f"attn_on_target.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                level="task",
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Attn on Target for Accuracy
            plotter.plot_correlation(
                x_data=evaluator.get_accuracies(as_lists=True),
                y_data=evaluator.attn_on_target.all,
                x_label="Accuracy",
                y_label="Attention on Target Tokens",
                file_name=f"acc-attn_on_target.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                level="task",
                include_soft=False,
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Correlation of accuracy with answer_not_mentioned and empty_attn_scores
            for extra_metric in ("answer_not_mentioned", "empty_attn_scores"):
                metric_obj = getattr(evaluator, extra_metric, None)
                if metric_obj is not None:
                    y_vals = (
                        metric_obj.all
                        if hasattr(metric_obj, "all")
                        else list(metric_obj)
                    )
                    plotter.plot_correlation(
                        x_data=evaluator.get_accuracies(as_lists=True),
                        y_data=y_vals,
                        x_label="Accuracy",
                        y_label=extra_metric.replace("_", " ").title(),
                        file_name=f"acc-{extra_metric}.pdf",
                        plot_name_add=[f"Task-{task_id}", *conditions_add],
                        path_add=Path(version, f"Task-{task_id}"),
                        level="task",
                        include_soft=False,
                        experiment=experiment,
                        num_samples=samples_per_task,
                    )

            # Attn on Target for Target Distances by Answer Correct
            plotter.plot_corr_boxplot(
                x_data=task.parts_target_distances.all,
                y_data={
                    "parts_attn_on_target": evaluator.parts_attn_on_target.all,
                    "parts_answer_correct": evaluator.parts_answer_correct.all,
                    "parts_features": task.parts_features[version],
                },
                x_label="Target Sentence Distances",
                y_label="Attention On Target",
                displ_percentage=False,
                version=version,
                file_name=f"attn-target_distances.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Attn on Target for Answer Correct by Parts Features
            plotter.plot_corr_boxplot(
                x_data=evaluator.parts_answer_correct.all,
                y_data={
                    "parts_attn_on_target": evaluator.parts_attn_on_target,
                    "parts_features": task.parts_features[version],
                },
                x_label="Answer Correct",
                y_label="Attention On Target",
                displ_percentage=False,
                version=version,
                file_name=f"attn-ans_correct.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Attn on target for Anwer in Self by Answer Correct
            plotter.plot_corr_boxplot(
                x_data=task.parts_answer_in_self.all,
                y_data={
                    "parts_attn_on_target": evaluator.parts_attn_on_target.all,
                    "parts_answer_correct": evaluator.parts_answer_correct.all,
                },
                x_label="Answer In Self",
                y_label="Attention On Target",
                displ_percentage=False,
                version=version,
                file_name=f"attn-ans_in_self.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Attn on Target for Seen Context Lengths by Answer Correct
            plotter.plot_corr_boxplot(
                x_data=task.seen_context_lengths.all,  # Added .all to convert to list/array for part-level plotting
                y_data={
                    "parts_attn_on_target": evaluator.parts_attn_on_target.all,
                    "parts_answer_correct": evaluator.parts_answer_correct.all,
                },
                x_label="Seen Context Lengths",
                y_label="Attention On Target",
                displ_percentage=False,
                version=version,
                file_name=f"attn-seen_context_lengths.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                experiment=experiment,
                num_samples=samples_per_task,
            )

            # Answer Correct for Seen Context Lengths by Answer In Self
            plotter.plot_corr_hist(
                x_data=task.seen_context_lengths.all,  # Added .all to convert to list/array for part-level plotting
                y_data={
                    "parts_answer_correct": evaluator.parts_answer_correct.all,
                    "parts_answer_in_self": task.parts_answer_in_self,
                },
                x_label="Seen Context Lengths",
                y_label="Parts Answer In[Correct]",
                displ_percentage=True,
                file_name=f"parts_answer_correct.pdf",
                plot_name_add=[f"Task-{task_id}", *conditions_add],
                path_add=Path(version, f"Task-{task_id}"),
                experiment=experiment,
                num_samples=samples_per_task,
            )
            plotter.correlation_map(
                data=corr_matrix,
                level=evaluator.level,
                version=version,
                file_name=f"corr_matrix_task_{task_id}.pdf",
                path_add=Path(version, f"Task-{task_id}"),
                id=task_id,
            )

            saver.save_json(
                data=corr_matrix,
                file_path=f"corr_matrix_task_{task_id}.json",
                path_add=Path(version, f"Task-{task_id}"),
            )

            metrics_to_save = defaultdict(dict)
            metrics = list(
                format_metrics(evaluator.get_metrics(as_lists=True)).values()
            )
            for metric in metrics:
                metrics_to_save[metric["task_id"]].update(metric)

            for metric in metrics_to_save.values():
                saver.save_output(
                    data=[metric],
                    headers=list(metric.keys()),
                    file_name=f"eval_script_metrics_{version}.csv",
                    path_add=Path(version),
                )

            print(
                f"\nPlotting distractor attention analysis for task {task_id} '{version}'...",
                end="\n\n",
            )
            d_stats_task = distractor_stats_per_task[task_id][version]
            if not d_stats_task.is_empty():
                plotter.plot_distractor_attn_boxplot(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_distractor_attn_per_task(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_distractor_attn_scatter(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_supporting_attention(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_distractor_supporting_ratio(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_attention_triplet(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_distraction_vs_n_distractors(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )
                plotter.plot_accuracy_vs_distraction_ratio(
                    stats=d_stats_task,
                    version=version,
                    plot_name_add=[f"Task-{task_id}", version, *conditions_add],
                    path_add=Path(version, f"Task-{task_id}"),
                )

                saver.save_output(
                    data=d_stats_task.as_csv_records(),
                    headers=d_stats_task.csv_headers,
                    file_name=f"distractor_attention_{version}.csv",
                    path_add=Path(version, f"Task-{task_id}"),
                )
            else:
                print(
                    f"No distractor attention records collected for task {task_id} version='{version}'. "
                    "Check that parts have interpretability data and distractors set."
                )

    if verbose:
        print_metrics_table(evaluators=split.evaluators, id_=data_split)

    saver.save_split_metrics(
        split=split,
        metric_file_name="eval_script_metrics.csv",
    )

    save_latex_table_line(split, experiment, setting, saver)

    split_corr_matrices = split.calculate_metrics()

    # ---- Before/after comparison plots --------------------------------------
    # These four plots compare the same metrics across the "before" and
    # "after" evaluators of this split. They run unconditionally — for
    # single-system runs the methods just plot the one available version.
    if len(split.evaluators) >= 1:
        print(
            f"\nPlotting before/after comparison plots for split '{split.name}'...",
            end="\n\n",
        )
        ba_kwargs = dict(
            evaluators=split.evaluators,
            versions=split.versions,
            plot_name_add=[f"Split-{split.name}", *conditions_add],
            path_add=Path("before_after"),
        )
        plotter.plot_before_after_accuracy(**ba_kwargs)
        if experiment != "direct_answer":
            plotter.plot_before_after_reasoning_scores(**ba_kwargs)
        plotter.plot_before_after_attention(**ba_kwargs)
        plotter.plot_before_after_summary(**ba_kwargs)

        if len(split.evaluators) > 1:
            plotter.plot_before_after_delta_lineplot(**ba_kwargs)

    for version, evaluator, features, corr_matrix in zip(
        split.versions, split.evaluators, split.features, split_corr_matrices.values()
    ):
        # SAVING
        saver.save_json(
            data=corr_matrix,
            file_path=f"corr_matrix_split_{split.name}.json",
            path_add=version,
        )
        saver.save_split_features(
            features=features,
            metrics_file_name="eval_script_features.csv",
            version=version,
        )

        # PLOTTING
        plotter.correlation_map(
            data=corr_matrix,
            level=evaluator.level,
            version=version,
            split_name=split.name,
            file_name=f"corr_matrix_split_{split.name}.pdf",
            path_add=Path(version),
        )

        # Plot Accuracy vs Attn on Target for the Split
        plotter.plot_correlation(
            x_data=evaluator.get_accuracies(as_lists=True),
            y_data=evaluator.attn_on_target.all,
            x_label="Accuracy",
            y_label="Attention on Target Tokens",
            file_name=f"acc-attn_on_target_{split.name}.pdf",
            plot_name_add=[f"Split-{split.name}", *conditions_add],
            experiment=experiment,
            num_samples=samples_per_task * len(split.tasks),
            path_add=Path(version),
            level="split",
            include_soft=False,
            label_add=[f"t{task.task_id}" for task in split.tasks],
        )

        # Correlation of accuracy with answer_not_mentioned and empty_attn_scores
        for extra_metric in ("answer_not_mentioned", "empty_attn_scores"):
            metric_obj = getattr(evaluator, extra_metric, None)
            if metric_obj is not None:
                y_vals = (
                    metric_obj.all if hasattr(metric_obj, "all") else list(metric_obj)
                )
                plotter.plot_correlation(
                    x_data=evaluator.get_accuracies(as_lists=True),
                    y_data=y_vals,
                    x_label="Accuracy",
                    y_label=extra_metric.replace("_", " ").title(),
                    file_name=f"acc-{extra_metric}_{split.name}.pdf",
                    plot_name_add=[f"Split-{split.name}", *conditions_add],
                    experiment=experiment,
                    num_samples=samples_per_task * len(split.tasks),
                    path_add=Path(version),
                    level="split",
                    include_soft=False,
                    label_add=[f"t{task.task_id}" for task in split.tasks],
                )

        # Attn on Target for Seen Context Lengths by Answer Correct
        plotter.plot_corr_boxplot(
            x_data=split.seen_context_lengths,
            y_data={
                "parts_attn_on_targets": evaluator.parts_attn_on_target.all,
                "parts_answer_correct": evaluator.parts_answer_correct.all,
            },
            x_label="Seen Context Lengths",
            y_label="Attention On Target",
            displ_percentage=False,
            version=version,
            experiment=experiment,
            num_samples=samples_per_task * len(split.tasks),
            level="split",
            file_name=f"attn-seen_context_lengths_{split.name}.pdf",
            plot_name_add=[f"Split-{split.name}", *conditions_add],
            path_add=Path(version),
        )

        # Attn on Target for Target Distances by Answer Correct
        plotter.plot_corr_boxplot(
            x_data=split.parts_target_distances,
            y_data={
                "parts_attn_on_targets": evaluator.parts_attn_on_target.all,
                "parts_answer_correct": evaluator.parts_answer_correct.all,
            },
            x_label="Target Sentence Distances",
            y_label="Attention On Target",
            level="split",
            displ_percentage=False,
            version=version,
            experiment=experiment,
            num_samples=samples_per_task * len(split.tasks),
            file_name=f"attn-target_distances_{split.name}.pdf",
            plot_name_add=[f"Split-{split.name}", *conditions_add],
            path_add=Path(version),
        )

        # Answer Correct for Seen Context Lengths by Answer In Self
        plotter.plot_corr_hist(
            x_data={"parts_seen_context_lengths": split.seen_context_lengths},
            y_data={
                "parts_answer_correct": evaluator.parts_answer_correct.all,
                "parts_answer_in_self": split.parts_answer_in_self,
            },
            x_label="Parts Seen Context Lengths",
            y_label="Parts Answer [In]Correct",
            level="split",
            displ_percentage=True,
            file_name=f"parts_answer_correct_{split.name}.pdf",
            plot_name_add=[f"Split-{split.name}", *conditions_add],
            path_add=Path(version),
            experiment=experiment,
            num_samples=samples_per_task * len(split.tasks),
        )
        print(
            f"\nPlotting accuracies and standard deviation for results '{version}'...",
            end="\n\n",
        )
        plotter.plot_acc_with_std(
            acc_per_prompt_task=evaluator.get_accuracies(as_lists=True),
            y_label="Accuracies with Standard Deviations",
            plot_name_add=[split.name, version, *conditions_add],
            path_add=Path(version),
        )
        print(
            f"\nPlotting attentions for results '{version}'...",
            end="\n\n",
        )
        plotter.plot_acc_with_std(
            acc_per_prompt_task=evaluator.get_attentions(as_lists=True),
            y_label="Attentions",
            plot_name_add=[split.name, version, *conditions_add],
            path_add=Path(version),
        )
        print(
            f"\nPlotting reasoning scores for results '{version}'...",
            end="\n\n",
        )
        if experiment != "direct_answer":
            plotter.plot_acc_with_std(
                acc_per_prompt_task=evaluator.get_reasoning_scores(as_lists=True),
                y_label="Reasoning Scores",
                plot_name_add=[split.name, version, *conditions_add],
                path_add=Path(version),
            )
        print(
            f"\nPlotting correlations for results '{version}' between metrics:",
            evaluator.get_correlations(as_lists=True),
            end="\n\n",
        )

        print(
            f"\nPlotting distractor attention analysis for '{version}'...", end="\n\n"
        )
        d_stats = distractor_stats[version]
        if not d_stats.is_empty():
            plotter.plot_distractor_attn_boxplot(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_distractor_attn_per_task(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_distractor_attn_scatter(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_supporting_attention(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_distractor_supporting_ratio(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_attention_triplet(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_distraction_vs_n_distractors(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )
            plotter.plot_accuracy_vs_distraction_ratio(
                stats=d_stats,
                version=version,
                plot_name_add=[f"Split-{split.name}", version, *conditions_add],
                path_add=Path(version),
            )

            saver.save_output(
                data=d_stats.as_csv_records(),
                headers=d_stats.csv_headers,
                file_name=f"distractor_attention_{version}.csv",
                path_add=Path(version),
            )
        else:
            print(
                f"No distractor attention records collected for version='{version}'. "
                "Check that parts have interpretability data and distractors set."
            )

        print("Saving result categories...")
        plotter.plot_answer_type_per_part(
            Results.CASE_COUNTERS[version],
            specification={
                "setting": setting,
                "experiment": experiment,
                "version": version,
            },
        )
        for score in ("bleu", "rouge", "meteor"):
            plotter.plot_answer_type_per_part(
                Results.CASE_COUNTERS[version],
                specification={
                    "setting": setting,
                    "experiment": experiment,
                    "version": version,
                    "score": score.upper(),
                },
                reasoning_scores=getattr(evaluator, f"ids_with_{score}"),
            )
        plotter.plot_answer_type_per_part(
            Results.CASE_COUNTERS[version],
            specification={
                "setting": setting,
                "experiment": experiment,
                "version": version,
                "score": "ATTN_ON_TARGET",
            },
            reasoning_scores=evaluator.ids_with_attn_on_target,
        )
        for case, case_list in Results.CASE_COUNTERS[version].items():
            headers = "id_\ttask_id\tsample_id\tpart_id"
            if case_list:
                saver.save_with_separator(
                    saver.run_path / version / f"{case}.txt",
                    [headers] + case_list,
                    sep="\n",
                )
                print(f"Case {case}: detected {len(case_list)} occurrences.")
            else:
                print(f"Case {case}: detected 0 occurrences. Nothing!")

    print(f"Plots produced: {plotter.plot_counter_prompt}")

    # Annotate the per-row results CSV with a completeness column so downstream
    # analyses can filter out rows where the model was truncated mid-reasoning
    # or never produced an answer. Both versions ("before" and "after") are
    # processed if their answer column exists; each writes a separate
    # _with_completeness_<version>.csv next to the original.
    results_csv_path = str(saver.run_path / results_file_name)
    if Path(results_csv_path).exists():
        try:
            df_for_check = pd.read_csv(results_csv_path, nrows=0)
        except Exception as err:
            warnings.warn(
                f"Could not read header of {results_csv_path} to add "
                f"completeness column: {err}"
            )
            df_for_check = None

        if df_for_check is not None:
            is_da = experiment == "direct_answer"
            for version in ("before", "after"):
                if f"model_answer_{version}" in df_for_check.columns:
                    add_completeness_column(
                        path=results_csv_path,
                        da=is_da,
                        before=(version == "before"),
                        max_tokens=max_tokens,
                        output_suffix=f"_{version}",
                    )
                    print(
                        f"Wrote completeness column for version='{version}' to "
                        f"{Path(results_csv_path).stem}_with_completeness_{version}.csv"
                    )
    else:
        warnings.warn(
            f"Per-row results CSV not found at {results_csv_path}; "
            "skipping completeness annotation."
        )

    print("\nThe evaluation pipeline has finished successfully.")


def parse_args(script_args: str | list[str] | None = None) -> argparse.Namespace:
    """
    Parse the command line arguments.

    :param script_args: optional list of argument strings; if None, sys.argv is used
    :return: parsed argument namespace
    """
    parser = argparse.ArgumentParser(description="Evaluate the results of the model.")
    parser.add_argument(
        "--results_path",
        type=str,
        required=True,
        help="Path to the data.",
    )
    parser.add_argument(
        "--save_path",
        type=str,
        required=True,
        help="Path where to save the results.",
    )
    parser.add_argument(
        "--setting",
        type=str,
        required=True,
        choices=[
            "baseline",
            "feedback",
            "skyline",
            "speculative_decoding",
            "sd",
            "basic-baseline",
        ],
        help="The setting of the experiment to evaluate.",
    )
    parser.add_argument(
        "--experiment",
        type=str,
        required=True,
        choices=["reasoning", "direct_answer"],
        help="The experiment to evaluate.",
    )
    parser.add_argument(
        "--samples_per_task",
        type=int,
        default=50,
        help="Number of samples per task the results were run with (check your config for the run).",
    )
    parser.add_argument(
        "--filtering_conditions",
        type=json.loads,
        default={},
        help=(
            "A JSON string representing a dictionary of conditions (SamplePart attributes) and values; "
            "all the parts having an attribute with such value will be preserved."
        ),
    )
    parser.add_argument(
        "--create_heatmaps",
        action="store_true",
        help="Whether to create attention heatmaps for the interpretability results.",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Whether to print the results to the console.",
    )
    parser.add_argument(
        "--max_tokens",
        type=int,
        default=None,
        help=(
            "The model's per-call generation budget. If a row's reasoning hits "
            "this many tokens it is flagged as truncated (incomplete) by the "
            "completeness annotation step. If omitted, the truncation check "
            "is skipped."
        ),
    )
    parser.add_argument(
        "--reasoning_source",
        type=str,
        default=None,
        choices=["claude", "llama"],
        help=(
            "Which silver-reasoning corpus to use when computing reasoning "
            "quality scores. "
            "'claude' reads from data/silver_reasoning/claude/; "
            "'llama' reads from data/silver_reasoning/llama/. "
            "Omit this argument to use the default flat directory "
            "data/silver_reasoning/ (legacy behaviour)."
        ),
    )
    if script_args is not None:
        if isinstance(script_args, str):
            script_args = script_args.split()
        args, unexpected = parser.parse_known_args(
            script_args, namespace=argparse.Namespace()
        )
        if unexpected:
            print(f"Unexpected arguments: {unexpected}")
        return args
    return parser.parse_args()


def add_completeness_column(
    path: str,
    da: bool = False,
    before: bool = False,
    max_tokens: int | None = None,
    reasoning_tokens_col: str | None = None,
    output_suffix: str = "",
) -> None:
    """
    Add a ``completeness`` column to a results CSV indicating whether the
    model's output for that row is complete.

    Criteria (non-DA case, i.e. reasoning + answer):
      - The answer is incomplete if no answer text was produced.
      - The reasoning is incomplete if either:
          (a) the answer is empty -- a chain of thought that never reaches an
              answer is, by definition, unfinished, or
          (b) the reasoning hit the exact token-generation limit, which means
              the model was cut off mid-thought rather than terminating
              naturally.
      - A row is complete iff both reasoning and answer are complete.

    Criteria (DA case, i.e. direct answer only, no reasoning):
      - Complete iff the answer is non-empty.

    The token-limit check (b) requires a pre-computed token count column.
    The function looks at ``reasoning_tokens_col`` (default
    ``f"model_reasoning_tokens_{suffix}"``). If ``max_tokens`` is None, or
    the token column is absent, the truncation check is skipped (a warning
    is emitted in the latter case) and only the "answer non-empty"
    criterion is applied to the reasoning.

    :param path: path to the results CSV to update
    :param da: whether the experiment is direct answer (no reasoning column)
    :param before: whether to check the "before" columns (default: "after")
    :param max_tokens: the generation budget the model was run with. If a
                       row's reasoning has exactly this many tokens the
                       reasoning is treated as truncated.
    :param reasoning_tokens_col: name of an existing column carrying the
                                  reasoning token count, if any
    :param output_suffix: extra string appended to the output filename. When
                          calling this for both "before" and "after" within
                          one run, pass "_before" / "_after" to keep the two
                          outputs distinct (default ``""`` keeps the
                          original behaviour).
    :return: None; writes ``<path>_with_completeness<output_suffix>.csv``
    """
    df = pd.read_csv(path)
    suffix = "before" if before else "after"

    def _is_nonempty(x) -> bool:
        return bool(pd.notna(x) and isinstance(x, str) and x.strip() != "")

    answer_col = f"model_answer_{suffix}"
    answer_complete = df[answer_col].apply(_is_nonempty)

    if da:
        df["completeness"] = answer_complete
        df.to_csv(
            f"{path.split('.')[0]}_with_completeness{output_suffix}.csv",
            index=False,
        )
        return

    reasoning_col = f"model_reasoning_{suffix}"
    reasoning_nonempty = df[reasoning_col].apply(_is_nonempty)

    if max_tokens is not None:
        token_col = reasoning_tokens_col or f"model_reasoning_tokens_{suffix}"
        if token_col in df.columns:
            reasoning_tokens = pd.to_numeric(df[token_col], errors="coerce")
            reasoning_truncated = reasoning_tokens >= max_tokens
            divisible_by_50 = reasoning_tokens % 50 == 0

            if divisible_by_50.any():
                multiple_of_50 = reasoning_tokens[divisible_by_50]
                warnings.warn(
                    "Some rows have reasoning token counts divisible by 50, "
                    "which may indicate generation limits were reached.\n"
                    f"Rows:\n{multiple_of_50}"
                )

        else:
            warnings.warn(
                f"add_completeness_column: column '{token_col}' not found; "
                "skipping the reasoning truncation check. To enable it, add a "
                "per-row reasoning token-count column before calling this "
                "function (or pass reasoning_tokens_col=...)."
            )

            reasoning_truncated = pd.Series(False, index=df.index)

    else:
        reasoning_truncated = pd.Series(False, index=df.index)

    reasoning_complete = reasoning_nonempty & ~reasoning_truncated & answer_complete

    df["completeness"] = reasoning_complete & answer_complete

    df.to_csv(
        f"{path.split('.')[0]}_with_completeness{output_suffix}.csv",
        index=False,
    )


def save_latex_table_line(split, experiment: str, setting: str, saver) -> None:
    """
    Write a ready-to-paste LaTeX row to latex_table_line.txt.
    ± is evaluator.X.get_std() at split level = std of per-task means across
    tasks (cross-task variability), which is the correct quantity for the
    paper table — distinct from the within-task std columns in the CSV.
    Table 1: EM-Acc | SM-Acc | Max Supp Attn | Attn on Trgt
    Table 2: BLEU | ROUGE | METEOR  (reasoning experiment only)

    :param split: the Split object containing the evaluators with the metrics to report
    :param experiment: the experiment name (e.g., "reasoning" or "direct_answer") to determine which metrics to include
    :param setting: the experimental setting (e.g., "baseline", "feedback", etc.) to include in the comment for context
    :param saver: the Saver object to use for writing the output file
    """
    _cell = lambda m, s: f"${round(m, 2)}^{{{{\\pm}}}}{round(s, 2)}$"
    _latex_lines: list[str] = []
    for _ver, _ev in zip(split.versions, split.evaluators):
        _em = _cell(
            _ev.exact_match_accuracy.get_mean(), _ev.exact_match_accuracy.get_std()
        )
        _sm = _cell(
            _ev.soft_match_accuracy.get_mean(), _ev.soft_match_accuracy.get_std()
        )
        _msa = _cell(_ev.max_supp_attn.get_mean(), _ev.max_supp_attn.get_std())
        _aot = _cell(_ev.attn_on_target.get_mean(), _ev.attn_on_target.get_std())
        _latex_lines.append(f"% Table 1 | {experiment} | {setting} | {_ver}")
        _latex_lines.append(f"& {_em} & {_sm} & {_msa} & {_aot} \\\\")
        if experiment == "reasoning":
            _b = _cell(_ev.bleu.get_mean(), _ev.bleu.get_std())
            _r = _cell(_ev.rouge.get_mean(), _ev.rouge.get_std())
            _m = _cell(_ev.meteor.get_mean(), _ev.meteor.get_std())
            _latex_lines.append(f"% Table 2 | {experiment} | {setting} | {_ver}")
            _latex_lines.append(f"& {_b} & {_r} & {_m} \\\\")
        _latex_lines.append("")
    (saver.run_path / "latex_table_line.txt").write_text(
        "\n".join(_latex_lines), encoding="utf-8"
    )
    print(f"LaTeX table line(s) written → {saver.run_path / 'latex_table_line.txt'}")


if __name__ == "__main__":
    args = parse_args()
    run(
        results_path=args.results_path,
        save_path=args.save_path,
        samples_per_task=args.samples_per_task,
        setting=args.setting,
        experiment=args.experiment,
        filtering_conditions={},
        create_heatmaps=args.create_heatmaps,
        verbose=args.verbose,
        max_tokens=args.max_tokens,
        reasoning_source=args.reasoning_source,
    )
    # kwargs = {
    #     "results_path": "/workspace/students/reasoning/results/basic-baseline/test/da/v1/all_tasks_joined/joined_direct_answer_results.csv",
    #     "save_path": "results/basic-baseline",
    #     "samples_per_task": 2,
    #     "setting": "baseline",
    #     "experiment": "direct_answer",
    #     "filtering_conditions": {},
    #     "create_heatmaps": True,
    #     "max_tokens": 1000, # or None
    #     "verbose": True,
    # }
    # print("Starting evaluation with the following configuration:")
    # print(json.dumps(kwargs, indent=2))
    # run(**kwargs)