mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
453 lines
14 KiB
Python
453 lines
14 KiB
Python
from dataclasses import dataclass
|
|
|
|
import matplotlib as mpl
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
# Styling (mirroring the other plot for visual consistency)
|
|
latte_style = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
|
|
plt.style.use(["ggplot", latte_style])
|
|
plt.rcParams["axes.prop_cycle"] = mpl.cycler(
|
|
color=[
|
|
"#FFB83D",
|
|
"#EEE684",
|
|
"#555555",
|
|
"#9BC750",
|
|
"#35775A",
|
|
"#5571EA",
|
|
"#6CACFF",
|
|
# "#E84494",
|
|
# "#E47FB0",
|
|
# "#785EF0",
|
|
# "#FF924E",
|
|
# "#5BC5DB",
|
|
"#AD6F50",
|
|
"#A1A9AD",
|
|
]
|
|
)
|
|
plt.rcParams["axes.facecolor"] = "white"
|
|
plt.rcParams["figure.facecolor"] = "white"
|
|
plt.rcParams["savefig.facecolor"] = "white"
|
|
|
|
# plt.rcParams["font.family"] = "Ubuntu"
|
|
# plt.rcParams["font.size"] = 14
|
|
# plt.rcParams["font.weight"] = "bold"
|
|
plt.rcParams["axes.labelweight"] = "bold"
|
|
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"] = 20
|
|
plt.rcParams["xtick.labelsize"] = 20
|
|
plt.rcParams["ytick.labelsize"] = 20
|
|
|
|
|
|
@dataclass
|
|
class Point:
|
|
performance: float
|
|
peak_mem: float | None = None # in GB
|
|
|
|
|
|
# Dataset focused only on memory vs performance (relative performance computed later)
|
|
datasets = {
|
|
"2WikiMultihopQA": {
|
|
# We keep Base model w/ context only as a performance reference (no point plotted if memory absent)
|
|
"Base model w/ context": [
|
|
Point(performance=0.242, peak_mem=2.099 * 1024),
|
|
],
|
|
"Base model w/ truncated context": [
|
|
Point(performance=0.3387, peak_mem=1.194 * 1024),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(performance=0.2138, peak_mem=65.162),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.305,
|
|
peak_mem=18.196,
|
|
),
|
|
],
|
|
"CD (5 generated queries)": [
|
|
Point(
|
|
performance=0.2385,
|
|
peak_mem=32.501,
|
|
),
|
|
],
|
|
"Ours": [
|
|
Point(
|
|
performance=0.2901,
|
|
peak_mem=35.777,
|
|
),
|
|
],
|
|
"Ours + truncated context": [
|
|
Point(
|
|
performance=0.3585,
|
|
peak_mem=1.367 * 1024,
|
|
)
|
|
],
|
|
# "Ours (8k chunks)": [
|
|
# Point(
|
|
# performance=0.2901,
|
|
# peak_mem=35.777,
|
|
# ),
|
|
# ],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.2474,
|
|
peak_mem=37.776,
|
|
),
|
|
],
|
|
"LLMLingua-2": [
|
|
# Point(performance=0.3141, peak_mem=1.111 * 1024),
|
|
Point(performance=0.3169, peak_mem=1.018 * 1024),
|
|
Point(performance=0.3558, peak_mem=807.366),
|
|
Point(performance=0.3591, peak_mem=592.584),
|
|
Point(performance=0.2767, peak_mem=305.298),
|
|
Point(performance=0.2343, peak_mem=158.760),
|
|
],
|
|
},
|
|
"MultiFieldQA": {
|
|
"Base model w/ context": [
|
|
Point(performance=0.303, peak_mem=1.211 * 1024),
|
|
],
|
|
"Base model w/ truncated context": [
|
|
Point(performance=0.3938, peak_mem=1.084 * 1024),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(performance=0.0948, peak_mem=65.853),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.4101,
|
|
peak_mem=7.836,
|
|
),
|
|
],
|
|
"CD (5 generated queries)": [
|
|
Point(
|
|
performance=0.165,
|
|
peak_mem=31.587,
|
|
),
|
|
],
|
|
"Ours": [
|
|
Point(
|
|
performance=0.1895,
|
|
peak_mem=35.769,
|
|
),
|
|
],
|
|
"Ours + truncated context": [
|
|
Point(
|
|
performance=0.4293,
|
|
peak_mem=1.233 * 1024,
|
|
),
|
|
],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.083,
|
|
peak_mem=38.907,
|
|
),
|
|
],
|
|
"LLMLingua-2": [
|
|
Point(performance=0.4222, peak_mem=909.663),
|
|
Point(performance=0.3898, peak_mem=716.981),
|
|
Point(performance=0.3773, peak_mem=490.478),
|
|
Point(performance=0.2885, peak_mem=256.169),
|
|
Point(performance=0.2298, peak_mem=146.394),
|
|
],
|
|
},
|
|
"QASPER": {
|
|
# We keep Base model w/ context only as a performance reference (no point plotted if memory absent)
|
|
"Base model w/ context": [
|
|
Point(performance=0.346, peak_mem=2.217 * 1024),
|
|
],
|
|
"Base model w/ truncated context": [
|
|
Point(performance=0.3839, peak_mem=788.697),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(performance=0.0859, peak_mem=64.989),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.3627,
|
|
peak_mem=19.611,
|
|
),
|
|
],
|
|
"CD (5 generated queries)": [
|
|
Point(
|
|
performance=0.1712,
|
|
peak_mem=36.767,
|
|
),
|
|
],
|
|
"Ours": [
|
|
Point(
|
|
performance=0.2051,
|
|
peak_mem=46.130,
|
|
),
|
|
],
|
|
"Ours + truncated context": [
|
|
Point(
|
|
performance=0.3751,
|
|
peak_mem=873.099,
|
|
),
|
|
],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.1259,
|
|
peak_mem=54.176,
|
|
),
|
|
],
|
|
"LLMLingua-2": [
|
|
# Point(performance=0.3734, peak_mem=710.176),
|
|
Point(performance=0.3849, peak_mem=647.385),
|
|
Point(performance=0.3838, peak_mem=487.781),
|
|
Point(performance=0.3328, peak_mem=386.447),
|
|
Point(performance=0.3329, peak_mem=219.062),
|
|
Point(performance=0.2614, peak_mem=132.576),
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
def build_color_map(datasets):
|
|
# Collect all unique methods in a deterministic order
|
|
methods = []
|
|
for data in datasets.values():
|
|
for m in data.keys():
|
|
if m not in methods:
|
|
methods.append(m)
|
|
|
|
# Sort methods to ensure consistent ordering across runs
|
|
# methods = sorted(methods)
|
|
|
|
palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
|
|
return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
|
|
|
|
|
|
def plot_memory_vs_performance(datasets):
|
|
n = len(datasets)
|
|
# If multiple datasets, arrange horizontally; otherwise single standard figure
|
|
if n == 1:
|
|
fig, single_ax = plt.subplots(1, 1, figsize=(7, 5))
|
|
axes = [single_ax]
|
|
else:
|
|
fig, axes_arr = plt.subplots(1, n, figsize=(7 * n, 5))
|
|
# Ensure axes is a simple list
|
|
axes = list(np.atleast_1d(axes_arr))
|
|
|
|
color_map = build_color_map(datasets)
|
|
global_methods = set()
|
|
ours_methods = {"Ours", "Ours + truncated context"} # added
|
|
annot_name = { # added
|
|
"Ours": "Ours",
|
|
"Ours + truncated context": "Ours+ctx",
|
|
}
|
|
|
|
highlight_label = "Efficient and Effective\nInternalization"
|
|
|
|
for idx, (dataset_name, methods) in enumerate(datasets.items()):
|
|
ax = axes[idx]
|
|
ref_perf = methods["Base model w/ truncated context"][0].performance
|
|
base_wo_ctx_perf_rel = (
|
|
methods["Base model w/o context"][0].performance / ref_perf
|
|
if "Base model w/o context" in methods
|
|
else None
|
|
)
|
|
|
|
all_mems = []
|
|
all_rel_perfs = []
|
|
ours_present = [ # added: detect both ours points for staggering
|
|
m
|
|
for m in ours_methods
|
|
if m in methods and any(p.peak_mem is not None for p in methods[m])
|
|
]
|
|
for method, points in methods.items():
|
|
mem_points = [p for p in points if p.peak_mem is not None]
|
|
if not mem_points:
|
|
continue # skip reference method if no memory value
|
|
xs = [p.peak_mem for p in mem_points]
|
|
ys = [p.performance / ref_perf for p in mem_points]
|
|
# marker mapping (oracle=diamond, ours=star, others=circle)
|
|
marker = (
|
|
"D"
|
|
if method == "CD (oracle)"
|
|
else ("*" if method in ours_methods else "o")
|
|
)
|
|
size = 20 if "Ours" in method else 12
|
|
ax.plot(
|
|
xs,
|
|
ys,
|
|
linestyle="-",
|
|
marker=marker,
|
|
linewidth=2,
|
|
markersize=size,
|
|
label=method,
|
|
color=color_map[method],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
alpha=0.9,
|
|
)
|
|
# Add annotation for Ours, stagger if both are present
|
|
if method in ours_methods:
|
|
yoff = (
|
|
(10 if method == "Ours" else -10) if len(ours_present) == 2 else 6
|
|
)
|
|
place_left = method == "Ours + truncated context" and dataset_name in {
|
|
"2WikiMultihopQA",
|
|
"MultiFieldQA",
|
|
}
|
|
if place_left:
|
|
yoff = 10 # move annotation above the marker for these two plots
|
|
xoff = -5 if place_left else 8
|
|
ha = "right" if place_left else "left"
|
|
ax.annotate(
|
|
annot_name[method],
|
|
(xs[-1], ys[-1]),
|
|
xytext=(xoff, yoff),
|
|
textcoords="offset points",
|
|
fontsize=12,
|
|
color=color_map[method],
|
|
fontweight="bold",
|
|
ha=ha,
|
|
bbox=None,
|
|
)
|
|
|
|
all_mems.extend(xs)
|
|
all_rel_perfs.extend(ys)
|
|
global_methods.add(method)
|
|
|
|
if all_mems:
|
|
max_mem = max(all_mems)
|
|
min_mem = min(all_mems)
|
|
span = (
|
|
(max_mem - min_mem)
|
|
if max_mem > min_mem
|
|
else (max_mem if max_mem != 0 else 1)
|
|
)
|
|
ax.set_xlim(5, max_mem * 1.25)
|
|
|
|
# Set individual y-axis limits for each subplot
|
|
if all_rel_perfs:
|
|
min_y = min(all_rel_perfs)
|
|
max_y = max(all_rel_perfs)
|
|
y_range = max_y - min_y
|
|
margin = y_range * 0.1 if y_range > 0 else 0.1
|
|
ax.set_ylim(min_y - margin, max_y + margin)
|
|
|
|
ax.set_xscale("log")
|
|
ax.set_xlabel(
|
|
"Additional Memory Needed\nfor Generation (MB)",
|
|
fontweight="bold",
|
|
fontsize=22,
|
|
)
|
|
if idx == 0: # Only show y-label on first plot
|
|
ax.set_ylabel("Normalized Performance", fontweight="bold", fontsize=22)
|
|
ax.grid(True, alpha=0.3)
|
|
ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
|
|
if base_wo_ctx_perf_rel is not None:
|
|
ax.axhline(
|
|
base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
|
|
)
|
|
|
|
ax.set_title(dataset_name, fontweight="bold", fontsize=30)
|
|
|
|
# Build grouped legends
|
|
in_context = [
|
|
m
|
|
for m in [
|
|
"Base model w/ context",
|
|
"Base model w/ truncated context",
|
|
"LLMLingua-2",
|
|
]
|
|
if m in global_methods
|
|
]
|
|
in_param_order = [
|
|
"CD (oracle)",
|
|
"CD (5 generated queries)",
|
|
"Ours",
|
|
"Ours + truncated context",
|
|
"T2L",
|
|
"Base model w/o context",
|
|
]
|
|
in_param = [m for m in in_param_order if m in global_methods]
|
|
|
|
def make_handles(names):
|
|
return [
|
|
mpl.lines.Line2D(
|
|
[],
|
|
[],
|
|
label=m,
|
|
marker=(
|
|
"D"
|
|
if m == "CD (oracle)"
|
|
else ("*" if m in {"Ours", "Ours + truncated context"} else "o")
|
|
),
|
|
linestyle="-",
|
|
linewidth=2,
|
|
color=color_map[m],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
markersize=22 if "Ours" in m else 14,
|
|
)
|
|
for m in names
|
|
]
|
|
|
|
ic_handles = make_handles(in_context)
|
|
ip_handles = make_handles(in_param)
|
|
|
|
# Extra bottom margin for legends
|
|
# (reduced wspace & hspace to tighten gaps)
|
|
# Increased hspace to add more vertical separation between dataset rows
|
|
plt.subplots_adjust(right=0.88, wspace=0.12, hspace=1.0, bottom=0.25)
|
|
|
|
# Set legend title fontweight before creating legends
|
|
# plt.rcParams["legend.title_fontsize"] = 14
|
|
|
|
fig.legend(
|
|
ic_handles,
|
|
[h.get_label() for h in ic_handles],
|
|
title="In-context knowledge",
|
|
loc="upper center",
|
|
bbox_to_anchor=(0.2, -0.05),
|
|
ncol=2,
|
|
frameon=True,
|
|
fancybox=True,
|
|
fontsize=18,
|
|
title_fontproperties={"weight": "bold", "size": 20},
|
|
)
|
|
fig.legend(
|
|
ip_handles,
|
|
[h.get_label() for h in ip_handles],
|
|
title="In-parameter knowledge",
|
|
loc="upper center",
|
|
bbox_to_anchor=(0.7, -0.05),
|
|
ncol=min(len(ip_handles), 3),
|
|
frameon=True,
|
|
fancybox=True,
|
|
fontsize=18,
|
|
title_fontproperties={"weight": "bold", "size": 20},
|
|
)
|
|
|
|
# Adjust spacing: horizontal layout uses wspace, single layout unaffected
|
|
if n == 1:
|
|
plt.subplots_adjust(bottom=0.18, left=0.12, right=0.98)
|
|
else:
|
|
# Tighter horizontal spacing and margins
|
|
plt.subplots_adjust(bottom=0.18, left=0.06, right=0.99, wspace=0.15)
|
|
return fig
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fig = plot_memory_vs_performance(datasets)
|
|
plt.show()
|
|
fig.savefig(
|
|
"/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.png",
|
|
dpi=300,
|
|
bbox_inches="tight",
|
|
)
|
|
fig.savefig(
|
|
"/home/tan/research/ctx-to-lora/tmp/performance_gen_mem.pdf",
|
|
dpi=300,
|
|
bbox_inches="tight",
|
|
)
|