mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
786 lines
26 KiB
Python
786 lines
26 KiB
Python
from dataclasses import dataclass # added
|
|
|
|
import matplotlib as mpl
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
# Styling (reused)
|
|
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",
|
|
"#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"] = 16
|
|
plt.rcParams["xtick.labelsize"] = 13
|
|
plt.rcParams["ytick.labelsize"] = 13
|
|
|
|
|
|
@dataclass
|
|
class Point:
|
|
performance: float
|
|
context_ratio: float | None = None
|
|
latency: float | None = None
|
|
latency_std: float | None = None
|
|
peak_mem: float | None = None # added
|
|
peak_mem_std: float | None = None # added
|
|
|
|
|
|
def combine_std_for_sum(std_list):
|
|
return np.sqrt(sum(std**2 for std in std_list))
|
|
|
|
|
|
# Unified data structure (merged context + latency via Point)
|
|
datasets = {
|
|
"SQuAD": {
|
|
"Base model w/ context": [
|
|
Point(performance=0.8692, context_ratio=1.0),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(
|
|
performance=0.1866,
|
|
context_ratio=0.0,
|
|
latency=0.0,
|
|
latency_std=0.0,
|
|
peak_mem=0.0,
|
|
),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.859,
|
|
context_ratio=0.0,
|
|
latency=39.974 + 456.989 * 1e-3,
|
|
latency_std=combine_std_for_sum([965.326 * 1e-3, 259.130 * 1e-3]),
|
|
peak_mem=855.762 / 1024,
|
|
),
|
|
],
|
|
"CD (generated queries)": [
|
|
Point(
|
|
performance=0.5986,
|
|
context_ratio=0.0,
|
|
latency=21.254 * 4 + 49.358 + 7.450,
|
|
latency_std=10,
|
|
peak_mem=41.098,
|
|
),
|
|
],
|
|
"Ours (batched)": [
|
|
Point(
|
|
# performance=0.7174,
|
|
performance=0.7077,
|
|
# context_ratio=0.0,
|
|
latency=85.954 * 1e-3 + 259.920 * 1e-6,
|
|
latency_std=combine_std_for_sum([31.057 * 1e-3, 740.286 * 1e-6]),
|
|
peak_mem=1.151,
|
|
),
|
|
],
|
|
"Ours (iterative)": [
|
|
Point(
|
|
performance=0.715,
|
|
context_ratio=0.0,
|
|
latency=426.849 * 1e-3 + 260.073 * 1e-6,
|
|
latency_std=combine_std_for_sum([43.423 * 1e-3, 29.673 * 1e-6]),
|
|
peak_mem=328.134 / 1024,
|
|
),
|
|
],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.1761,
|
|
context_ratio=0.0,
|
|
latency=35.480 * 1e-3,
|
|
latency_std=22.193 * 1e-3,
|
|
peak_mem=1.793,
|
|
),
|
|
],
|
|
"LLMLingua-2": [
|
|
Point(performance=0.8492, context_ratio=0.9),
|
|
Point(performance=0.8278, context_ratio=0.8),
|
|
Point(performance=0.7749, context_ratio=0.6),
|
|
Point(performance=0.6624, context_ratio=0.4),
|
|
Point(performance=0.4271, context_ratio=0.2),
|
|
Point(performance=0.3192, context_ratio=0.1),
|
|
],
|
|
},
|
|
"DROP": {
|
|
"Base model w/ context": [
|
|
Point(performance=0.4541, context_ratio=1.0),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(
|
|
performance=0.1417,
|
|
context_ratio=0.0,
|
|
latency=0.0,
|
|
latency_std=0.0,
|
|
peak_mem=0.0,
|
|
),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.443,
|
|
context_ratio=0.0,
|
|
latency=39.974 + 422.820 * 1e-3,
|
|
latency_std=combine_std_for_sum([965.326 * 1e-3, 241.419 * 1e-3]),
|
|
peak_mem=1.381,
|
|
),
|
|
],
|
|
"CD (generated queries)": [
|
|
Point(
|
|
performance=0.2289,
|
|
context_ratio=0.0,
|
|
latency=23.883 * 4 + 9.755 + 41.301,
|
|
latency_std=14.8852904910,
|
|
peak_mem=44.305,
|
|
),
|
|
],
|
|
"Ours (batched)": [
|
|
Point(
|
|
performance=0.2972,
|
|
# context_ratio=0.0,
|
|
latency=84.132 * 1e-3 + 251.876 * 1e-6,
|
|
latency_std=combine_std_for_sum([30.503 * 1e-3, 22.829 * 1e-6]),
|
|
peak_mem=1.976,
|
|
),
|
|
],
|
|
"Ours (iterative)": [
|
|
Point(
|
|
# performance=0.3002,
|
|
performance=0.3387,
|
|
context_ratio=0.0,
|
|
latency=419.462 * 1e-3 + 254.899 * 1e-6,
|
|
latency_std=combine_std_for_sum([40.211 * 1e-3, 29.784 * 1e-6]),
|
|
peak_mem=429.881 / 1024,
|
|
),
|
|
],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.1468,
|
|
context_ratio=0.0,
|
|
latency=35.036 * 1e-3,
|
|
latency_std=18.144 * 1e-3,
|
|
peak_mem=3.004,
|
|
),
|
|
],
|
|
"LLMLingua-2": [ #
|
|
Point(performance=0.4513, context_ratio=0.9),
|
|
Point(performance=0.4488, context_ratio=0.8),
|
|
Point(performance=0.4076, context_ratio=0.6),
|
|
Point(performance=0.3488, context_ratio=0.4),
|
|
Point(performance=0.2617, context_ratio=0.2),
|
|
Point(performance=0.1712, context_ratio=0.1),
|
|
],
|
|
},
|
|
"ROPES": {
|
|
"Base model w/ context": [
|
|
Point(performance=0.7457, context_ratio=1.0),
|
|
],
|
|
"Base model w/o context": [
|
|
Point(
|
|
performance=0.4545,
|
|
context_ratio=0.0,
|
|
latency=0.0,
|
|
latency_std=0.0,
|
|
peak_mem=0.0,
|
|
),
|
|
],
|
|
"CD (oracle)": [
|
|
Point(
|
|
performance=0.7349,
|
|
context_ratio=0.0,
|
|
latency=39.633 + 285.467 * 1e-3,
|
|
latency_std=combine_std_for_sum([1.752, 72.791 * 1e-3]),
|
|
peak_mem=527.312 / 1024,
|
|
),
|
|
],
|
|
"CD (generated queries)": [
|
|
Point(
|
|
performance=0.5786,
|
|
context_ratio=0.0,
|
|
latency=23.714 * 4 + 11.473 + 59.660,
|
|
latency_std=11.7329040736,
|
|
peak_mem=43.271,
|
|
),
|
|
],
|
|
"Ours (batched)": [
|
|
Point(
|
|
performance=0.6759,
|
|
# context_ratio=0.0,
|
|
latency=85.652 * 1e-3 + 254.029 * 1e-6,
|
|
latency_std=combine_std_for_sum([30.852 * 1e-3, 51.127 * 1e-6]),
|
|
peak_mem=496.898 / 1024,
|
|
),
|
|
],
|
|
"Ours (iterative)": [
|
|
Point(
|
|
# performance=0.6786,
|
|
performance=0.6088,
|
|
context_ratio=0.0,
|
|
latency=417.860 * 1e-3 + 257.422 * 1e-6,
|
|
latency_std=combine_std_for_sum([45.514 * 1e-3, 31.419 * 1e-6]),
|
|
peak_mem=273.867 / 1024,
|
|
),
|
|
],
|
|
"T2L": [
|
|
Point(
|
|
performance=0.4661,
|
|
context_ratio=0.0,
|
|
latency=37.913 * 1e-3,
|
|
latency_std=17.313 * 1e-3,
|
|
peak_mem=697.217 / 1024,
|
|
),
|
|
],
|
|
"LLMLingua-2": [
|
|
Point(performance=0.708, context_ratio=0.9),
|
|
Point(performance=0.7126, context_ratio=0.8),
|
|
Point(performance=0.7291, context_ratio=0.6),
|
|
Point(performance=0.677, context_ratio=0.4),
|
|
Point(performance=0.6188, context_ratio=0.2),
|
|
Point(performance=0.4894, context_ratio=0.1),
|
|
],
|
|
},
|
|
}
|
|
|
|
|
|
def build_color_map(datasets):
|
|
methods = []
|
|
for data in datasets.values():
|
|
for m in data.keys():
|
|
if m not in methods:
|
|
methods.append(m)
|
|
palette = plt.rcParams["axes.prop_cycle"].by_key()["color"]
|
|
return {m: palette[i % len(palette)] for i, m in enumerate(methods)}
|
|
|
|
|
|
def plot_ctx_and_latency(datasets):
|
|
n = len(datasets)
|
|
# detect memory usage presence
|
|
memory_present = any(
|
|
p.peak_mem is not None
|
|
for methods in datasets.values()
|
|
for plist in methods.values()
|
|
for p in plist
|
|
)
|
|
num_cols = 3 if memory_present else 2
|
|
# Share x-axes across rows only when there are multiple datasets (n > 1)
|
|
if n > 1:
|
|
fig, axes = plt.subplots(
|
|
n,
|
|
num_cols,
|
|
figsize=(6 * num_cols, 5 * n),
|
|
sharey=True,
|
|
sharex="col",
|
|
)
|
|
else:
|
|
fig, axes = plt.subplots(
|
|
n, num_cols, figsize=(6 * num_cols, 5 * n), sharey=True
|
|
)
|
|
if n == 1:
|
|
axes = np.array([axes])
|
|
|
|
color_map = build_color_map(datasets)
|
|
global_methods = set()
|
|
|
|
label_text = "Efficient and Effective\nInternalization"
|
|
ours_methods = {"Ours (batched)", "Ours (iterative)"} # added
|
|
annot_name = {
|
|
"Ours (batched)": "Ours (b)",
|
|
"Ours (iterative)": "Ours (i)",
|
|
}
|
|
|
|
# Keep track of axes groups per row for later separator lines (always store as list)
|
|
row_axes_groups = []
|
|
for row_idx, (dataset_name, methods) in enumerate(datasets.items()):
|
|
row_axes = axes[row_idx]
|
|
if memory_present:
|
|
ctx_ax, lat_ax, mem_ax = row_axes
|
|
else:
|
|
ctx_ax, lat_ax = row_axes
|
|
# Ensure each row group is a plain Python list of Axes objects
|
|
if isinstance(row_axes, (list, tuple)):
|
|
row_axes_groups.append(list(row_axes))
|
|
elif isinstance(row_axes, np.ndarray):
|
|
row_axes_groups.append(list(row_axes.ravel()))
|
|
else: # single axis fallback (should not occur with current layout)
|
|
row_axes_groups.append([row_axes])
|
|
# Row subtitle (dataset name) centered across the row using the middle axis
|
|
title_ax = lat_ax # use latency axis (middle) for centering
|
|
title_ax.set_title(dataset_name, fontweight="bold", fontsize=24, pad=20)
|
|
|
|
# Reference performance from ICL
|
|
ref_perf = methods["Base model w/ context"][0].performance
|
|
base_wo_ctx_perf_rel = (
|
|
next(
|
|
p.performance
|
|
for p in methods["Base model w/o context"]
|
|
if p.context_ratio is not None
|
|
)
|
|
/ ref_perf
|
|
)
|
|
|
|
# --- Context subplot ---
|
|
ours_ctx_present = [ # added
|
|
m
|
|
for m in ours_methods
|
|
if any(p.context_ratio is not None for p in methods.get(m, []))
|
|
]
|
|
for method, points in methods.items():
|
|
ctx_points = [p for p in points if p.context_ratio is not None]
|
|
if not ctx_points:
|
|
continue
|
|
xs = [p.context_ratio for p in ctx_points]
|
|
ys = [p.performance / ref_perf for p in ctx_points]
|
|
marker = (
|
|
"D"
|
|
if method == "CD (oracle)"
|
|
else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o")
|
|
)
|
|
markersize = 17 if "Ours" in method else 12
|
|
|
|
if len(xs) == 1:
|
|
ctx_ax.scatter(
|
|
xs,
|
|
ys,
|
|
marker=marker,
|
|
s=markersize**2, # scatter uses area, so square the size
|
|
label=method,
|
|
color=color_map[method],
|
|
edgecolors="black",
|
|
linewidths=1.5,
|
|
alpha=0.9,
|
|
)
|
|
else:
|
|
ctx_ax.plot(
|
|
xs,
|
|
ys,
|
|
marker=marker,
|
|
linewidth=2,
|
|
markersize=markersize,
|
|
label=method,
|
|
color=color_map[method],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
alpha=0.9,
|
|
)
|
|
if method in {"Ours (batched)", "Ours (iterative)"}:
|
|
yoff = (
|
|
(10 if method == "Ours (batched)" else -10)
|
|
if len(ours_ctx_present) == 2
|
|
else 6
|
|
)
|
|
|
|
ctx_ax.annotate( # fixed to ctx_ax
|
|
annot_name[method],
|
|
(xs[-1], ys[-1]),
|
|
xytext=(8, yoff),
|
|
textcoords="offset points",
|
|
fontsize=12,
|
|
color=color_map[method],
|
|
fontweight="bold",
|
|
bbox=None,
|
|
)
|
|
global_methods.add(method)
|
|
|
|
# Add x-axis label on every context subplot (reverted per-user preference)
|
|
ctx_ax.set_xlabel("Context Length Ratio", fontweight="bold", fontsize=16)
|
|
# if row_idx == 0:
|
|
ctx_ax.set_ylabel(
|
|
"Normalized Performance",
|
|
fontweight="bold",
|
|
fontsize=16,
|
|
)
|
|
ctx_ax.grid(True, alpha=0.3)
|
|
ctx_ax.set_xlim(-0.05, 1.05)
|
|
ctx_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
|
|
ctx_ax.axhline(
|
|
base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
|
|
)
|
|
# Removed highlight rectangle and label for context subplot
|
|
# --- Latency subplot ---
|
|
ours_lat_present = [ # added
|
|
m
|
|
for m in ours_methods
|
|
if any(p.latency is not None for p in methods.get(m, []))
|
|
]
|
|
all_latencies = []
|
|
for method, points in methods.items():
|
|
lat_points = [p for p in points if p.latency is not None]
|
|
if not lat_points:
|
|
continue
|
|
lat_x = [p.latency for p in lat_points]
|
|
perf_rel = [p.performance / ref_perf for p in lat_points]
|
|
xerr = [p.latency_std for p in lat_points]
|
|
marker = (
|
|
"D"
|
|
if method == "CD (oracle)"
|
|
else ("*" if method in {"Ours (batched)", "Ours (iterative)"} else "o")
|
|
)
|
|
markersize = 17 if "Ours" in method else 12
|
|
|
|
if len(lat_x) == 1 and all(err == 0 for err in xerr):
|
|
lat_ax.scatter(
|
|
lat_x,
|
|
perf_rel,
|
|
marker=marker,
|
|
s=markersize**2,
|
|
label=method,
|
|
color=color_map[method],
|
|
edgecolors="black",
|
|
linewidths=1.5,
|
|
alpha=0.9,
|
|
)
|
|
else:
|
|
lat_ax.errorbar(
|
|
lat_x,
|
|
perf_rel,
|
|
xerr=xerr,
|
|
fmt=marker + "-",
|
|
linewidth=2,
|
|
markersize=markersize,
|
|
capsize=4,
|
|
capthick=1,
|
|
label=method,
|
|
color=color_map[method],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
alpha=0.9,
|
|
)
|
|
# Add floating annotation for Ours
|
|
if method in {"Ours (batched)", "Ours (iterative)"}:
|
|
yoff = (
|
|
(-20 if method == "Ours (batched)" else 10)
|
|
if len(ours_lat_present) == 2
|
|
else 6
|
|
)
|
|
lat_ax.annotate(
|
|
annot_name[method],
|
|
(lat_x[-1], perf_rel[-1]),
|
|
xytext=(-25, yoff),
|
|
textcoords="offset points",
|
|
fontsize=12,
|
|
color=color_map[method],
|
|
fontweight="bold",
|
|
bbox=None,
|
|
)
|
|
|
|
all_latencies.extend(lat_x)
|
|
global_methods.add(method)
|
|
lat_ax.set_xscale("log")
|
|
if all_latencies:
|
|
positive = [x for x in all_latencies if x > 0]
|
|
if positive:
|
|
min_pos = min(positive)
|
|
max_x = max(all_latencies)
|
|
lat_ax.set_xlim(min_pos * 0.8, max_x * 1.2)
|
|
lat_ax.set_xlabel("Update Latency (seconds)", fontweight="bold", fontsize=16)
|
|
lat_ax.grid(True, which="both", alpha=0.3)
|
|
lat_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
|
|
lat_ax.axhline(
|
|
base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
|
|
)
|
|
# Add vertical line at 1 second
|
|
lat_ax.axvline(1.0, color="gray", linestyle="--", alpha=0.7, zorder=-1)
|
|
lat_ax.text(
|
|
1.1,
|
|
0.5,
|
|
"Sub-second internalization",
|
|
transform=lat_ax.get_xaxis_transform(),
|
|
ha="left",
|
|
va="center",
|
|
fontsize=8,
|
|
color="gray",
|
|
fontstyle="italic",
|
|
rotation=90,
|
|
)
|
|
# # Highlight rectangle
|
|
# mpl.patches.Rectangle(
|
|
# (0.0, 0.5),
|
|
# 0.5,
|
|
# 0.5,
|
|
# transform=lat_ax.transAxes,
|
|
# facecolor="green",
|
|
# edgecolor="none",
|
|
# alpha=0.15,
|
|
# zorder=-2,
|
|
# )
|
|
# )
|
|
# lat_ax.text(
|
|
# 0.02,
|
|
# 0.62,
|
|
# label_text,
|
|
# transform=lat_ax.transAxes,
|
|
# fontsize=8.5,
|
|
# ha="left",
|
|
# va="bottom",
|
|
# color="green",
|
|
# fontweight="bold",
|
|
# )
|
|
|
|
# --- Memory subplot (new) ---
|
|
if memory_present:
|
|
ours_mem_present = [ # added
|
|
m
|
|
for m in ours_methods
|
|
if any(p.peak_mem is not None for p in methods.get(m, []))
|
|
]
|
|
all_mems = []
|
|
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
|
|
mem_x = [p.peak_mem for p in mem_points]
|
|
perf_rel = [p.performance / ref_perf for p in mem_points]
|
|
xerr = [
|
|
p.peak_mem_std if p.peak_mem_std is not None else 0
|
|
for p in mem_points
|
|
]
|
|
marker = (
|
|
"D"
|
|
if method == "CD (oracle)"
|
|
else (
|
|
"*" if method in {"Ours (batched)", "Ours (iterative)"} else "o"
|
|
)
|
|
)
|
|
markersize = 17 if "Ours" in method else 12
|
|
|
|
if len(mem_x) == 1 and all(err == 0 for err in xerr):
|
|
mem_ax.scatter(
|
|
mem_x,
|
|
perf_rel,
|
|
marker=marker,
|
|
s=markersize**2,
|
|
label=method,
|
|
color=color_map[method],
|
|
edgecolors="black",
|
|
linewidths=1.5,
|
|
alpha=0.9,
|
|
)
|
|
else:
|
|
mem_ax.errorbar(
|
|
mem_x,
|
|
perf_rel,
|
|
xerr=xerr,
|
|
fmt=marker + "-",
|
|
linewidth=2,
|
|
markersize=markersize,
|
|
capsize=4,
|
|
capthick=1,
|
|
label=method,
|
|
color=color_map[method],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
alpha=0.9,
|
|
)
|
|
# Add floating annotation for Ours
|
|
if method in {"Ours (batched)", "Ours (iterative)"}:
|
|
yoff = (
|
|
(-20 if method == "Ours (batched)" else 15)
|
|
if len(ours_mem_present) == 2
|
|
else 6
|
|
)
|
|
mem_ax.annotate(
|
|
annot_name[method],
|
|
(mem_x[-1], perf_rel[-1]),
|
|
xytext=(-15, yoff),
|
|
textcoords="offset points",
|
|
fontsize=12,
|
|
color=color_map[method],
|
|
fontweight="bold",
|
|
bbox=None,
|
|
)
|
|
|
|
all_mems.extend(mem_x)
|
|
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
|
|
)
|
|
mem_ax.set_xlim(min_mem - 0.05 * span, max_mem * 1.05)
|
|
mem_ax.set_xlabel(
|
|
"Additional Memory Needed\nfor Model Updates (GB)",
|
|
fontweight="bold",
|
|
fontsize=16,
|
|
)
|
|
mem_ax.grid(True, alpha=0.3)
|
|
mem_ax.axhline(1.0, color="gray", linestyle="--", alpha=0.5, zorder=-1)
|
|
mem_ax.axhline(
|
|
base_wo_ctx_perf_rel, color="gray", linestyle="--", alpha=0.5, zorder=-1
|
|
)
|
|
# # Highlight rectangle (memory) fixed 20% width in axes coords
|
|
# mem_ax.add_patch(
|
|
# mpl.patches.Rectangle(
|
|
# (0.0, 0.6),
|
|
# 0.4,
|
|
# 0.4,
|
|
# transform=mem_ax.transAxes,
|
|
# facecolor="green",
|
|
# edgecolor="none",
|
|
# alpha=0.15,
|
|
# zorder=-2,
|
|
# )
|
|
# )
|
|
# mem_ax.text(
|
|
# 0.02,
|
|
# 0.62,
|
|
# label_text,
|
|
# transform=mem_ax.transAxes,
|
|
# fontsize=8.5,
|
|
# ha="left",
|
|
# va="bottom",
|
|
# color="green",
|
|
# fontweight="bold",
|
|
# )
|
|
|
|
# REMOVE old single legend block
|
|
# (deleted fig.legend(...) that was on the right)
|
|
|
|
# Build grouped legends
|
|
in_context = [
|
|
m for m in ["Base model w/ context", "LLMLingua-2"] if m in global_methods
|
|
]
|
|
in_param_order = [
|
|
"CD (oracle)",
|
|
"CD (generated queries)",
|
|
"Ours (iterative)",
|
|
"Ours (batched)",
|
|
"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 (batched)", "Ours (iterative)"} else "o")
|
|
),
|
|
linestyle="-",
|
|
linewidth=2,
|
|
color=plt.rcParams["axes.prop_cycle"].by_key()["color"][
|
|
list(build_color_map(datasets).keys()).index(m)
|
|
% len(plt.rcParams["axes.prop_cycle"].by_key()["color"])
|
|
],
|
|
markeredgecolor="black",
|
|
markeredgewidth=1.5,
|
|
markersize=20 if "Ours" in m else 10,
|
|
)
|
|
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=0.8, bottom=0.25)
|
|
|
|
# Re-enable x tick labels for all rows (matplotlib hides them for shared axes not on bottom)
|
|
if n > 1:
|
|
for ax in axes.reshape(-1):
|
|
ax.tick_params(labelbottom=True)
|
|
|
|
# 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.25, 0.1),
|
|
ncol=len(ic_handles),
|
|
frameon=True,
|
|
fancybox=True,
|
|
fontsize=12,
|
|
title_fontproperties={"weight": "bold", "size": 14},
|
|
)
|
|
fig.legend(
|
|
ip_handles,
|
|
[h.get_label() for h in ip_handles],
|
|
title="In-parameter knowledge",
|
|
loc="upper center",
|
|
bbox_to_anchor=(0.6, 0.1),
|
|
ncol=min(len(ip_handles), 3),
|
|
frameon=True,
|
|
fancybox=True,
|
|
fontsize=12,
|
|
title_fontproperties={"weight": "bold", "size": 14},
|
|
)
|
|
|
|
# Removed global suptitle referencing only the last dataset; per-row titles now provided.
|
|
|
|
# Add horizontal separator lines between dataset rows (if more than one dataset)
|
|
if len(row_axes_groups) > 1:
|
|
# Use figure coordinate system
|
|
for upper_group, lower_group in zip(row_axes_groups[:-1], row_axes_groups[1:]):
|
|
# Reference axes (first axis in each group)
|
|
upper_ax = upper_group[0]
|
|
lower_ax = lower_group[0]
|
|
upper_bottom = upper_ax.get_position().y0
|
|
lower_top = lower_ax.get_position().y1
|
|
line_y = (upper_bottom + lower_top) / 2.0
|
|
combined = upper_group + lower_group
|
|
x0 = min(ax.get_position().x0 for ax in combined)
|
|
x1 = max(ax.get_position().x1 for ax in combined)
|
|
fig.add_artist(
|
|
mpl.lines.Line2D(
|
|
[x0, x1],
|
|
[line_y, line_y],
|
|
transform=fig.transFigure,
|
|
color="#cccccc",
|
|
linewidth=1.2,
|
|
alpha=0.8,
|
|
)
|
|
)
|
|
return fig
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fig = plot_ctx_and_latency(datasets)
|
|
plt.show()
|
|
fig.savefig(
|
|
"/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.png",
|
|
dpi=300,
|
|
bbox_inches="tight",
|
|
)
|
|
fig.savefig(
|
|
"/home/tan/research/ctx-to-lora/tmp/performance_ctx_latency_mem.pdf",
|
|
dpi=300,
|
|
bbox_inches="tight",
|
|
)
|