plots.py

import os
import json
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import re
import argparse
import seaborn as sns
import sys

def record_judge_data(root, dirs, files, prompt_type=None):
    """
    Record data from json file given the directory. Changes based on prompt type.

    :param root:
    :param dirs:
    :param files:
    :param prompt_type:
    :return: records: Dict of the recorded data
    """

    assert prompt_type in ("reasoning", "multi-hop")

    # Parameters
    judge_filename = 'judged_result.json'

    if prompt_type == 'reasoning':
        prompt_index_pattern = r"_p[1245]$"
        whole_reasoning_chain_cols = [
            "makes_sense",
            "content_incomplete",
            "contradictory_steps",
            "irrelevant_info",
            "grammar_issue",
            "factual_or_math_error",
        ]
    else:
        prompt_index_pattern = r"_p[3678]$"
        whole_reasoning_chain_cols = [
            "makes_sense",
            "content_incomplete",
            "contradictory_steps",
        ]



    if judge_filename in files:
        fp = os.path.join(root, judge_filename)
        records = []
        with open(fp, "r", encoding="utf-8") as f:
            data = json.load(f)

        for entry in data:
            if not re.search(prompt_index_pattern, entry['index']):
                continue

            jr = entry["judge_response"]
            rec = {
                "model": root,
                "overall_answer_quality":
                    jr["whole_response"]["overall_answer_quality"],
            }
            chain = jr["whole_reasoning_chain"]

            for k in whole_reasoning_chain_cols:
                try:
                    rec[k] = chain[k]
                except:
                    pass

            if prompt_type == 'multi-hop':
                try:
                    for key, value in jr.items():
                        if not key.startswith("hop"):
                            continue

                        for subkey, subval in value.items():
                            rec[f"{key}_{subkey}"] = subval
                except:
                    pass
            records.append(rec)

        return records


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Plotting script for Judge Metrics")
    parser.add_argument("--prompt", type=str, required=True, help="Reasoning vs. multi-hop selection")
    parser.add_argument("--directory", type=str, required=False, nargs='+', help="Specific directories containing judge results")
    args = parser.parse_args()

    prompt_type = args.prompt.lower().strip()
    assert prompt_type in ("reasoning", "multi-hop")

    if prompt_type == 'reasoning':
        whole_reasoning_chain_cols = [
            "overall_answer_quality",
            "makes_sense",
            "content_incomplete",
            "contradictory_steps",
            "irrelevant_info",
            "grammar_issue",
            "factual_or_math_error",
        ]
    else:
        whole_reasoning_chain_cols = [
            "makes_sense",
            "content_incomplete",
            "contradictory_steps",
        ]
        hop_cols = [
            "hop_type_correct",
            "redundant_hop",
            "grammar_issue",
            "context_contradiction",
            "logical_error",
            "wrong_final_answer",
            "irrelevant_info",
            "unnecessary_factual_info",
            "repetition",
            "general_knowledge_error",
            "math_error"
        ]

    records = []
    for root, dirs, files in os.walk("."):
        dirname = os.path.basename(root)
        if args.directory is not None and dirname not in args.directory:
            continue

        record = record_judge_data(root, dirs, files, prompt_type=prompt_type)
        records.append((record, dirname))

    df_list = [(pd.DataFrame(record), dirname) for record, dirname in records]
    
    # Plot analysis on whole reasoning chain
    for df, model in df_list:
        long_df = df[whole_reasoning_chain_cols].melt(var_name="metric", value_name="score")
        # long_df["score"] = long_df["score"].astype(int)

        palette = sns.color_palette("viridis", 6)  # 5-color palette
        sns.countplot(data=long_df, x="metric", hue="score", palette=palette)
        plt.xticks(rotation=45)
        plt.title(f"Count Plot on Reasoning Only, {model}")
        plt.tight_layout()

        output_folder = "plots"
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        filename = f"{model}_countplot.png"
        full_path = os.path.join(output_folder, filename)
        plt.savefig(full_path)
        plt.clf()  # Clear figure for next plot

    # Plot analysis on hop-by-hop basis
    if prompt_type != 'multi-hop':
        sys.exit(0)

    for hop_num in range(1, 5):  # 1-4 hops in our dataset
        hop_pattern = re.compile(rf"^hop{hop_num}_(.+)$")

        for col in df.columns:
            hop_data_cols = [col for col in df.columns if hop_pattern.match(col)]
            df_hops = df[hop_data_cols]

        fig, axes = plt.subplots(1, 2, figsize=(10, 8))
        fig.suptitle(f'Hop {hop_num}')

        sns.boxplot(ax=axes[0], data=df_hops[hop_cols])
        plt.xticks(rotation=45)
        plt.title("Quartile Distribution")
        # plt.show()

        sns.swarmplot(ax=axes[1], data=df_hops[hop_cols], size=4)
        plt.xticks(rotation=45)
        plt.title("Beeswarm Distribution")
        plt.show()