mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
277 lines
8.7 KiB
Python
277 lines
8.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
import matplotlib as mpl
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
from plot_ablation_utils import (
|
|
choose_latest_runs_by_dataset_and_chunk,
|
|
dataset_context_length,
|
|
ensure_output_path,
|
|
get_dense_values,
|
|
get_performance_value,
|
|
load_script_text,
|
|
parse_checkpoint_path,
|
|
parse_eval_results_root,
|
|
parse_expected_datasets,
|
|
)
|
|
|
|
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
|
|
|
|
|
|
def configure_plot_style() -> None:
|
|
plt.style.use(["ggplot", LATTE_STYLE])
|
|
plt.rcParams["axes.facecolor"] = "F7FBFC"
|
|
plt.rcParams["figure.facecolor"] = "white"
|
|
plt.rcParams["savefig.facecolor"] = "white"
|
|
plt.rcParams["grid.color"] = "cccccc"
|
|
plt.rcParams["grid.linewidth"] = 1
|
|
plt.rcParams["axes.edgecolor"] = "ccd0da"
|
|
plt.rcParams["legend.facecolor"] = "white"
|
|
plt.rcParams["legend.fontsize"] = 14
|
|
plt.rcParams["xtick.labelsize"] = 13
|
|
plt.rcParams["ytick.labelsize"] = 13
|
|
plt.rcParams["axes.labelweight"] = "bold"
|
|
plt.rcParams["axes.prop_cycle"] = mpl.cycler(color=plt.cm.tab10.colors)
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="Plot the num-chunks ablation from 4-eval-vary-num-chunks.sh.",
|
|
)
|
|
parser.add_argument(
|
|
"--eval-script",
|
|
type=Path,
|
|
default=Path(__file__).with_name("4-eval-vary-num-chunks.sh"),
|
|
help="Path to the eval sweep shell script.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
default=None,
|
|
help="Optional output PNG path. Defaults to eval-results-*/plots/4-vary-num-chunks.png.",
|
|
)
|
|
parser.add_argument(
|
|
"--metric",
|
|
default="rougeL.f1",
|
|
help="Performance metric to read from each *_results.json file.",
|
|
)
|
|
parser.add_argument(
|
|
"--dense-metric",
|
|
default="mean",
|
|
help="Dense LoRA metric to read when available.",
|
|
)
|
|
parser.add_argument(
|
|
"--dense-std-metric",
|
|
default="std",
|
|
help="Dense LoRA std metric to use for the shaded area when available.",
|
|
)
|
|
parser.add_argument(
|
|
"--dense-max-metric",
|
|
default="max",
|
|
help="Dense LoRA max metric to plot as a separate line when available.",
|
|
)
|
|
return parser
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
configure_plot_style()
|
|
|
|
script_text = load_script_text(args.eval_script)
|
|
checkpoint_path = parse_checkpoint_path(script_text)
|
|
eval_results_root = parse_eval_results_root(checkpoint_path)
|
|
expected_datasets = parse_expected_datasets(script_text)
|
|
selected_runs_by_dataset = choose_latest_runs_by_dataset_and_chunk(
|
|
eval_results_root,
|
|
expected_datasets,
|
|
)
|
|
|
|
if not any(selected_runs_by_dataset.values()):
|
|
raise SystemExit(f"No matching eval runs found under {eval_results_root}.")
|
|
|
|
output_path = (
|
|
eval_results_root / "plots" / "4-vary-num-chunks.png"
|
|
if args.output is None
|
|
else args.output
|
|
)
|
|
output_path = ensure_output_path(output_path)
|
|
|
|
fig, (ax_perf, ax_dense, ax_dense_max) = plt.subplots(
|
|
1,
|
|
3,
|
|
figsize=(20, 5.5),
|
|
constrained_layout=False,
|
|
)
|
|
fig.subplots_adjust(bottom=0.27, top=0.88, wspace=0.28)
|
|
|
|
color_map = plt.get_cmap("tab10", len(expected_datasets))
|
|
all_num_chunks: set[int] = set()
|
|
has_any_dense_mean_values = False
|
|
has_any_dense_max_values = False
|
|
|
|
for index, dataset_name in enumerate(expected_datasets):
|
|
color = color_map(index)
|
|
context_length = dataset_context_length(dataset_name)
|
|
label = rf"$2^{{{int(context_length).bit_length() - 1}}}$"
|
|
|
|
num_chunks: list[int] = []
|
|
performance_values: list[float] = []
|
|
dense_mean_xs: list[int] = []
|
|
dense_mean_ys: list[float] = []
|
|
dense_max_xs: list[int] = []
|
|
dense_max_ys: list[float] = []
|
|
|
|
for run_info in selected_runs_by_dataset.get(dataset_name, []):
|
|
if run_info.chunk_size == 24:
|
|
continue
|
|
|
|
approx_num_chunks = max(1, context_length // run_info.chunk_size)
|
|
performance = get_performance_value(
|
|
run_info.run_dir,
|
|
dataset_name,
|
|
metric_name=args.metric,
|
|
)
|
|
if performance is None:
|
|
continue
|
|
|
|
num_chunks.append(approx_num_chunks)
|
|
performance_values.append(performance)
|
|
all_num_chunks.add(approx_num_chunks)
|
|
|
|
dense_stats = get_dense_values(
|
|
run_info.run_dir,
|
|
dataset_name,
|
|
metric_names=(
|
|
args.dense_metric,
|
|
args.dense_std_metric,
|
|
args.dense_max_metric,
|
|
),
|
|
)
|
|
dense_mean = dense_stats[args.dense_metric]
|
|
dense_max = dense_stats[args.dense_max_metric]
|
|
if dense_mean is not None:
|
|
dense_mean_xs.append(approx_num_chunks)
|
|
dense_mean_ys.append(dense_mean)
|
|
if dense_max is not None:
|
|
dense_max_xs.append(approx_num_chunks)
|
|
dense_max_ys.append(dense_max)
|
|
|
|
if performance_values:
|
|
ax_perf.plot(
|
|
num_chunks,
|
|
performance_values,
|
|
marker="o",
|
|
linewidth=2,
|
|
color=color,
|
|
label=label,
|
|
markeredgecolor="black",
|
|
markeredgewidth=1,
|
|
)
|
|
if dense_mean_ys:
|
|
has_any_dense_mean_values = True
|
|
ax_dense.plot(
|
|
dense_mean_xs,
|
|
dense_mean_ys,
|
|
marker="o",
|
|
linewidth=2,
|
|
color=color,
|
|
label=label,
|
|
markeredgecolor="black",
|
|
markeredgewidth=1,
|
|
)
|
|
if dense_max_ys:
|
|
has_any_dense_max_values = True
|
|
ax_dense_max.plot(
|
|
dense_max_xs,
|
|
dense_max_ys,
|
|
marker="o",
|
|
linewidth=2,
|
|
color=color,
|
|
label=label,
|
|
markeredgecolor="black",
|
|
markeredgewidth=1,
|
|
)
|
|
|
|
if not all_num_chunks:
|
|
raise SystemExit("No performance values were found for the selected runs.")
|
|
|
|
sorted_num_chunks = sorted(all_num_chunks)
|
|
for axis in (ax_perf, ax_dense, ax_dense_max):
|
|
axis.set_xscale("log", base=2)
|
|
axis.set_xticks(sorted_num_chunks)
|
|
axis.set_xticklabels(
|
|
[
|
|
rf"$2^{{{math.log2(value):.1f}}}$"
|
|
if not math.log2(value).is_integer()
|
|
else rf"$2^{{{int(math.log2(value))}}}$"
|
|
for value in sorted_num_chunks
|
|
],
|
|
rotation=0,
|
|
ha="center",
|
|
)
|
|
axis.grid(True, alpha=0.25)
|
|
axis.xaxis.set_minor_locator(plt.NullLocator())
|
|
|
|
ax_perf.set_title("Retrieval Performance", fontweight="bold")
|
|
ax_perf.set_xlabel("Number of Chunks")
|
|
ax_perf.set_ylabel("ROUGE-L F1 Score")
|
|
ax_perf.set_ylim(-0.05, 1.05)
|
|
|
|
ax_dense.set_title("LoRA Mean Magnitude", fontweight="bold")
|
|
ax_dense.set_xlabel("Number of Chunks")
|
|
ax_dense.set_ylabel("Mean Per-Element Magnitude")
|
|
if has_any_dense_mean_values:
|
|
ax_dense.set_ylim(bottom=0.0)
|
|
else:
|
|
ax_dense.text(
|
|
0.5,
|
|
0.5,
|
|
"Dense LoRA stats are missing for these runs.",
|
|
ha="center",
|
|
va="center",
|
|
transform=ax_dense.transAxes,
|
|
)
|
|
|
|
ax_dense_max.set_title("LoRA Max Magnitude", fontweight="bold")
|
|
ax_dense_max.set_xlabel("Number of Chunks")
|
|
ax_dense_max.set_ylabel("Maximum Per-Element Magnitude")
|
|
if has_any_dense_max_values:
|
|
ax_dense_max.set_ylim(bottom=0.0)
|
|
else:
|
|
ax_dense_max.text(
|
|
0.5,
|
|
0.5,
|
|
"Dense LoRA max stats are missing for these runs.",
|
|
ha="center",
|
|
va="center",
|
|
transform=ax_dense_max.transAxes,
|
|
)
|
|
|
|
handles, labels = ax_perf.get_legend_handles_labels()
|
|
haystack_legend = fig.legend(
|
|
handles,
|
|
labels,
|
|
loc="lower center",
|
|
bbox_to_anchor=(0.5, -0.005),
|
|
ncol=min(6, len(labels)),
|
|
frameon=True,
|
|
fancybox=True,
|
|
title="Haystack length",
|
|
title_fontsize=16,
|
|
)
|
|
haystack_legend.get_frame().set_edgecolor("#bdbdbd")
|
|
haystack_legend.get_frame().set_linewidth(1.0)
|
|
haystack_legend.get_frame().set_alpha(1.0)
|
|
|
|
fig.savefig(output_path, dpi=200, bbox_inches="tight")
|
|
print(f"Saved plot to {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|