mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
740 lines
23 KiB
Python
740 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import re
|
|
import shlex
|
|
from collections import defaultdict, deque
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
import matplotlib as mpl
|
|
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
|
|
LATTE_STYLE = "https://raw.githubusercontent.com/51616/catppuccin-matplotlib/main/src/mplcatppuccin/data/latte.mplstyle"
|
|
BASE_MODEL_EVAL_ROOT = Path("eval_results/google/gemma-2-2b-it")
|
|
DATASET_LABELS = {
|
|
"oolong-synth": "Oolong-Synth",
|
|
"longbench/qasper_e": "Qasper",
|
|
"longbench/2wikimqa_e": "2WikiMQA",
|
|
"longbench/multifieldqa_en_e": "MultiFieldQA-EN",
|
|
"longbench/gov_report_e": "GovReport",
|
|
}
|
|
METRIC_CANDIDATES = (
|
|
"oolong_score",
|
|
"qa_f1_score",
|
|
"rougeL.f1",
|
|
"qa_f1",
|
|
"f1",
|
|
"accuracy",
|
|
)
|
|
METRIC_LABELS = {
|
|
"oolong_score": "Oolong Score",
|
|
"qa_f1_score": "QA F1 Score",
|
|
"rougeL.f1": "ROUGE-L F1",
|
|
"qa_f1": "QA F1",
|
|
"f1": "F1",
|
|
"accuracy": "Accuracy",
|
|
}
|
|
COMPARISON_TOLERANCE = 0.2
|
|
CMD_RE = re.compile(r"CMD:\s+(.*)$", re.MULTILINE)
|
|
BASELINE_EXCLUDE_FLAGS = {
|
|
"--use_rag",
|
|
"--use_hybrid_rag",
|
|
"--use_cd",
|
|
"--use_llmlingua",
|
|
"--use_t2l",
|
|
}
|
|
NO_CONTEXT_COLUMN = "No context"
|
|
CSV_SUBCOLUMNS = (NO_CONTEXT_COLUMN, "top-1", "top-4")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RagRunInfo:
|
|
run_dir: Path
|
|
run_name: str
|
|
datasets: tuple[str, ...]
|
|
is_hybrid: bool
|
|
top_k: int
|
|
chunk_size: int
|
|
chunk_overlap: int
|
|
max_retrieved_tokens: int
|
|
|
|
|
|
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"] = 12
|
|
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 one raw-RAG vs hybrid-RAG comparison figure per dataset.",
|
|
)
|
|
parser.add_argument(
|
|
"--rag-root",
|
|
type=Path,
|
|
default=Path("eval_results/google/gemma-2-2b-it"),
|
|
help="Directory containing raw RAG eval runs.",
|
|
)
|
|
parser.add_argument(
|
|
"--hybrid-root",
|
|
type=Path,
|
|
default=Path("trained_d2l/gemma_2b_d2l/eval-results-20000"),
|
|
help="Directory containing hybrid RAG eval runs.",
|
|
)
|
|
parser.add_argument(
|
|
"--rag-chunk-size",
|
|
type=int,
|
|
default=256,
|
|
help="Filter runs by rag chunk size.",
|
|
)
|
|
parser.add_argument(
|
|
"--rag-chunk-overlap",
|
|
type=int,
|
|
default=64,
|
|
help="Filter runs by rag chunk overlap.",
|
|
)
|
|
parser.add_argument(
|
|
"--rag-max-retrieved-tokens",
|
|
type=int,
|
|
default=1536,
|
|
help="Filter runs by rag max retrieved tokens.",
|
|
)
|
|
parser.add_argument(
|
|
"--datasets",
|
|
nargs="*",
|
|
default=None,
|
|
help="Optional dataset allowlist. Defaults to all datasets with matched raw/hybrid runs.",
|
|
)
|
|
parser.add_argument(
|
|
"--top-ks",
|
|
nargs="*",
|
|
type=int,
|
|
default=None,
|
|
help="Optional top-k allowlist. Defaults to all matched top-k values.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=Path,
|
|
default=None,
|
|
help="Optional output directory. Defaults to HYBRID_ROOT/plots.",
|
|
)
|
|
return parser
|
|
|
|
|
|
def _extract_flag_value(tokens: list[str], flag: str) -> str | None:
|
|
for index, token in enumerate(tokens):
|
|
if token == flag and index + 1 < len(tokens):
|
|
return tokens[index + 1]
|
|
if token.startswith(f"{flag}="):
|
|
return token.split("=", 1)[1]
|
|
return None
|
|
|
|
|
|
def _extract_datasets(tokens: list[str]) -> tuple[str, ...]:
|
|
if "--datasets" not in tokens:
|
|
return ()
|
|
|
|
index = tokens.index("--datasets") + 1
|
|
datasets: list[str] = []
|
|
while index < len(tokens) and not tokens[index].startswith("--"):
|
|
datasets.append(tokens[index])
|
|
index += 1
|
|
return tuple(datasets)
|
|
|
|
|
|
def _has_max_test_samples_flag(tokens: list[str]) -> bool:
|
|
return _extract_flag_value(tokens, "--max_test_samples_per_ds") is not None
|
|
|
|
|
|
def parse_rag_run_info(run_dir: Path) -> RagRunInfo | None:
|
|
log_path = run_dir / "debug.log"
|
|
if not log_path.exists():
|
|
return None
|
|
|
|
log_text = log_path.read_text(encoding="utf-8", errors="replace")
|
|
match = CMD_RE.search(log_text)
|
|
if match is None:
|
|
return None
|
|
|
|
command = match.group(1).strip()
|
|
try:
|
|
tokens = shlex.split(command)
|
|
except ValueError:
|
|
tokens = command.split()
|
|
|
|
use_rag = "--use_rag" in tokens
|
|
use_hybrid_rag = "--use_hybrid_rag" in tokens
|
|
if use_rag == use_hybrid_rag:
|
|
return None
|
|
if _has_max_test_samples_flag(tokens):
|
|
return None
|
|
|
|
datasets = _extract_datasets(tokens)
|
|
if not datasets:
|
|
return None
|
|
|
|
top_k_value = _extract_flag_value(tokens, "--rag_top_k")
|
|
chunk_size_value = _extract_flag_value(tokens, "--rag_chunk_size")
|
|
chunk_overlap_value = _extract_flag_value(tokens, "--rag_chunk_overlap")
|
|
max_tokens_value = _extract_flag_value(tokens, "--rag_max_retrieved_tokens")
|
|
if None in (
|
|
top_k_value,
|
|
chunk_size_value,
|
|
chunk_overlap_value,
|
|
max_tokens_value,
|
|
):
|
|
return None
|
|
|
|
return RagRunInfo(
|
|
run_dir=run_dir,
|
|
run_name=run_dir.name,
|
|
datasets=datasets,
|
|
is_hybrid=use_hybrid_rag,
|
|
top_k=int(top_k_value),
|
|
chunk_size=int(chunk_size_value),
|
|
chunk_overlap=int(chunk_overlap_value),
|
|
max_retrieved_tokens=int(max_tokens_value),
|
|
)
|
|
|
|
|
|
def parse_plain_base_run_datasets(run_dir: Path) -> tuple[str, ...] | None:
|
|
log_path = run_dir / "debug.log"
|
|
if not log_path.exists():
|
|
return None
|
|
|
|
log_text = log_path.read_text(encoding="utf-8", errors="replace")
|
|
match = CMD_RE.search(log_text)
|
|
if match is None:
|
|
return None
|
|
|
|
command = match.group(1).strip()
|
|
try:
|
|
tokens = shlex.split(command)
|
|
except ValueError:
|
|
tokens = command.split()
|
|
|
|
if "--model_name_or_path" not in tokens:
|
|
return None
|
|
if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
|
|
return None
|
|
if _has_max_test_samples_flag(tokens):
|
|
return None
|
|
|
|
datasets = _extract_datasets(tokens)
|
|
if not datasets:
|
|
return None
|
|
return datasets
|
|
|
|
|
|
def parse_plain_d2l_run_datasets(run_dir: Path) -> tuple[str, ...] | None:
|
|
log_path = run_dir / "debug.log"
|
|
if not log_path.exists():
|
|
return None
|
|
|
|
log_text = log_path.read_text(encoding="utf-8", errors="replace")
|
|
match = CMD_RE.search(log_text)
|
|
if match is None:
|
|
return None
|
|
|
|
command = match.group(1).strip()
|
|
try:
|
|
tokens = shlex.split(command)
|
|
except ValueError:
|
|
tokens = command.split()
|
|
|
|
if "--checkpoint_path" not in tokens:
|
|
return None
|
|
if any(flag in tokens for flag in BASELINE_EXCLUDE_FLAGS):
|
|
return None
|
|
if "--add_ctx_to_input" in tokens:
|
|
return None
|
|
if _has_max_test_samples_flag(tokens):
|
|
return None
|
|
max_ctx_chunk_len = _extract_flag_value(tokens, "--max_ctx_chunk_len")
|
|
if max_ctx_chunk_len != "8192":
|
|
return None
|
|
|
|
datasets = _extract_datasets(tokens)
|
|
if not datasets:
|
|
return None
|
|
return datasets
|
|
|
|
|
|
def result_json_path(run_dir: Path, dataset_name: str) -> Path:
|
|
return run_dir / f"test_{dataset_name}_results.json"
|
|
|
|
|
|
def load_results_json(run_dir: Path, dataset_name: str) -> dict[str, object]:
|
|
return json.loads(
|
|
result_json_path(run_dir, dataset_name).read_text(encoding="utf-8")
|
|
)
|
|
|
|
|
|
def infer_metric_name(result_json: dict[str, object], dataset_name: str) -> str:
|
|
for metric_name in METRIC_CANDIDATES:
|
|
if f"test_{dataset_name}_{metric_name}" in result_json:
|
|
return metric_name
|
|
raise ValueError(f"Could not infer a metric for dataset {dataset_name}.")
|
|
|
|
|
|
def metric_label(metric_name: str) -> str:
|
|
return METRIC_LABELS.get(metric_name, metric_name.replace("_", " ").title())
|
|
|
|
|
|
def dataset_label(dataset_name: str) -> str:
|
|
return DATASET_LABELS.get(dataset_name, dataset_name.split("/")[-1])
|
|
|
|
|
|
def select_latest_runs(
|
|
root: Path,
|
|
*,
|
|
want_hybrid: bool,
|
|
allowed_datasets: set[str] | None,
|
|
allowed_top_ks: set[int] | None,
|
|
rag_chunk_size: int,
|
|
rag_chunk_overlap: int,
|
|
rag_max_retrieved_tokens: int,
|
|
) -> dict[tuple[str, int], RagRunInfo]:
|
|
selected: dict[tuple[str, int], RagRunInfo] = {}
|
|
if not root.exists():
|
|
return selected
|
|
|
|
for run_dir in sorted(root.iterdir()):
|
|
if not run_dir.is_dir():
|
|
continue
|
|
|
|
run_info = parse_rag_run_info(run_dir)
|
|
if run_info is None or run_info.is_hybrid != want_hybrid:
|
|
continue
|
|
if run_info.chunk_size != rag_chunk_size:
|
|
continue
|
|
if run_info.chunk_overlap != rag_chunk_overlap:
|
|
continue
|
|
if run_info.max_retrieved_tokens != rag_max_retrieved_tokens:
|
|
continue
|
|
if allowed_top_ks is not None and run_info.top_k not in allowed_top_ks:
|
|
continue
|
|
|
|
for dataset_name in run_info.datasets:
|
|
if allowed_datasets is not None and dataset_name not in allowed_datasets:
|
|
continue
|
|
if not result_json_path(run_info.run_dir, dataset_name).exists():
|
|
continue
|
|
|
|
key = (dataset_name, run_info.top_k)
|
|
current = selected.get(key)
|
|
if current is None or run_info.run_name > current.run_name:
|
|
selected[key] = run_info
|
|
|
|
return selected
|
|
|
|
|
|
def get_metric_value(
|
|
run_dir: Path, dataset_name: str, metric_name: str
|
|
) -> float | None:
|
|
result_json = load_results_json(run_dir, dataset_name)
|
|
value = result_json.get(f"test_{dataset_name}_{metric_name}")
|
|
if value in (None, "None", "N/A"):
|
|
return None
|
|
if isinstance(value, (int, float)):
|
|
return float(value)
|
|
return None
|
|
|
|
|
|
def select_latest_plain_base_runs(
|
|
root: Path,
|
|
datasets: set[str],
|
|
) -> dict[str, Path]:
|
|
selected: dict[str, Path] = {}
|
|
if not root.exists():
|
|
return selected
|
|
|
|
for run_dir in sorted(root.iterdir()):
|
|
if not run_dir.is_dir():
|
|
continue
|
|
|
|
run_datasets = parse_plain_base_run_datasets(run_dir)
|
|
if run_datasets is None:
|
|
continue
|
|
|
|
for dataset_name in run_datasets:
|
|
if dataset_name not in datasets:
|
|
continue
|
|
if not result_json_path(run_dir, dataset_name).exists():
|
|
continue
|
|
selected[dataset_name] = run_dir
|
|
|
|
return selected
|
|
|
|
|
|
def select_latest_plain_d2l_runs(
|
|
root: Path,
|
|
datasets: set[str],
|
|
) -> dict[str, Path]:
|
|
selected: dict[str, Path] = {}
|
|
if not root.exists():
|
|
return selected
|
|
|
|
for run_dir in sorted(root.iterdir()):
|
|
if not run_dir.is_dir():
|
|
continue
|
|
|
|
run_datasets = parse_plain_d2l_run_datasets(run_dir)
|
|
if run_datasets is None:
|
|
continue
|
|
|
|
for dataset_name in run_datasets:
|
|
if dataset_name not in datasets:
|
|
continue
|
|
if not result_json_path(run_dir, dataset_name).exists():
|
|
continue
|
|
selected[dataset_name] = run_dir
|
|
|
|
return selected
|
|
|
|
|
|
CSV_TOP_KS = (1, 4)
|
|
|
|
|
|
def csv_output_path(output_dir: Path) -> Path:
|
|
return output_dir / "d2l-rag-vs-hybrid.csv"
|
|
|
|
|
|
def confusion_output_path_for_dataset(output_dir: Path, dataset_name: str) -> Path:
|
|
return output_dir / f"d2l-rag-vs-hybrid-confusion-{dataset_name.split('/')[-1]}.png"
|
|
|
|
|
|
def generated_text_path(run_dir: Path, dataset_name: str) -> Path:
|
|
return run_dir / f"test_{dataset_name}_generated_text.jsonl"
|
|
|
|
|
|
def sample_key(sample: dict[str, object], dataset_name: str) -> str:
|
|
if dataset_name == "oolong-synth":
|
|
return "oolong:" + str(sample.get("oolong_id"))
|
|
|
|
input_text = str(sample.get("input", ""))
|
|
label_text = str(sample.get("label", ""))
|
|
ctx_ids_len = str(sample.get("ctx_ids_len", ""))
|
|
return f"{dataset_name}\n{input_text}\n{label_text}\n{ctx_ids_len}"
|
|
|
|
|
|
def sample_metric_value(sample: dict[str, object], metric_name: str) -> float:
|
|
value = sample.get(metric_name)
|
|
if not isinstance(value, (int, float)):
|
|
return 0.0
|
|
return float(value)
|
|
|
|
|
|
def load_metric_by_sample(
|
|
run_dir: Path,
|
|
dataset_name: str,
|
|
metric_name: str,
|
|
) -> dict[str, deque[float]]:
|
|
path = generated_text_path(run_dir, dataset_name)
|
|
grouped: dict[str, deque[float]] = defaultdict(deque)
|
|
with path.open(encoding="utf-8") as handle:
|
|
for line in handle:
|
|
sample = json.loads(line)
|
|
grouped[sample_key(sample, dataset_name)].append(
|
|
sample_metric_value(sample, metric_name)
|
|
)
|
|
return grouped
|
|
|
|
|
|
def relative_bucket(score: float, other_score: float) -> int:
|
|
delta = score - other_score
|
|
if delta > COMPARISON_TOLERANCE:
|
|
return 2
|
|
if delta < -COMPARISON_TOLERANCE:
|
|
return 0
|
|
return 1
|
|
|
|
|
|
def build_confusion_counts(
|
|
raw_run_dir: Path,
|
|
hybrid_run_dir: Path,
|
|
dataset_name: str,
|
|
metric_name: str,
|
|
) -> list[list[int]]:
|
|
raw_grouped = load_metric_by_sample(raw_run_dir, dataset_name, metric_name)
|
|
hybrid_grouped = load_metric_by_sample(hybrid_run_dir, dataset_name, metric_name)
|
|
raw_counts = {key: len(values) for key, values in raw_grouped.items()}
|
|
hybrid_counts = {key: len(values) for key, values in hybrid_grouped.items()}
|
|
if raw_counts != hybrid_counts:
|
|
raise ValueError(f"Could not align raw and hybrid samples for {dataset_name}.")
|
|
|
|
matrix = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
|
|
for key in sorted(raw_grouped):
|
|
raw_values = raw_grouped[key]
|
|
hybrid_values = hybrid_grouped[key]
|
|
while raw_values:
|
|
raw_score = raw_values.popleft()
|
|
hybrid_score = hybrid_values.popleft()
|
|
raw_bucket = relative_bucket(raw_score, hybrid_score)
|
|
hybrid_bucket = relative_bucket(hybrid_score, raw_score)
|
|
matrix[raw_bucket][hybrid_bucket] += 1
|
|
return matrix
|
|
|
|
|
|
def plot_confusion_matrices(
|
|
output_dir: Path,
|
|
dataset_name: str,
|
|
metric_name: str,
|
|
paired_points: list[tuple[int, float, float]],
|
|
raw_runs: dict[tuple[str, int], RagRunInfo],
|
|
hybrid_runs: dict[tuple[str, int], RagRunInfo],
|
|
) -> None:
|
|
ncols = len(paired_points)
|
|
fig, axes = plt.subplots(
|
|
1,
|
|
ncols,
|
|
figsize=(4.8 * ncols, 4.6),
|
|
constrained_layout=False,
|
|
)
|
|
if ncols == 1:
|
|
axes = [axes]
|
|
fig.subplots_adjust(bottom=0.16, top=0.80, wspace=0.35)
|
|
|
|
for ax, (top_k, _, _) in zip(axes, paired_points):
|
|
matrix = build_confusion_counts(
|
|
raw_runs[(dataset_name, top_k)].run_dir,
|
|
hybrid_runs[(dataset_name, top_k)].run_dir,
|
|
dataset_name,
|
|
metric_name,
|
|
)
|
|
image = ax.imshow(matrix, cmap="Blues")
|
|
total = sum(sum(row) for row in matrix)
|
|
max_cell = max(max(row) for row in matrix) if total else 0
|
|
|
|
for raw_index in range(3):
|
|
for hybrid_index in range(3):
|
|
count = matrix[raw_index][hybrid_index]
|
|
pct = 100.0 * count / total if total else 0.0
|
|
color = "white" if count > max_cell / 2 else "#4c4f69"
|
|
ax.text(
|
|
hybrid_index,
|
|
raw_index,
|
|
f"{count}\n({pct:.1f}%)",
|
|
ha="center",
|
|
va="center",
|
|
color=color,
|
|
fontsize=11,
|
|
fontweight="bold" if count == max_cell and count > 0 else None,
|
|
)
|
|
|
|
ax.set_xticks([0, 1, 2])
|
|
ax.set_yticks([0, 1, 2])
|
|
ax.set_xticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
|
|
ax.set_yticklabels(["More\nIncorrect", "Same", "More\nCorrect"])
|
|
ax.set_xlabel("Hybrid RAG")
|
|
if ax is axes[0]:
|
|
ax.set_ylabel("Raw RAG")
|
|
ax.set_title(f"top-k = {top_k}", fontweight="bold")
|
|
ax.grid(False)
|
|
|
|
for spine in ax.spines.values():
|
|
spine.set_visible(True)
|
|
spine.set_color("#ccd0da")
|
|
|
|
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
|
|
|
|
fig.suptitle(
|
|
f"Relative Score Agreement\n({dataset_label(dataset_name)}, delta threshold = {COMPARISON_TOLERANCE:.1f})",
|
|
fontweight="bold",
|
|
fontsize=16,
|
|
)
|
|
output_path = confusion_output_path_for_dataset(output_dir, dataset_name)
|
|
fig.savefig(output_path, dpi=200, bbox_inches="tight")
|
|
print(f"Saved plot to {output_path}")
|
|
plt.close(fig)
|
|
|
|
|
|
def write_summary_csv(
|
|
output_dir: Path,
|
|
datasets: list[str],
|
|
summary: dict[str, dict[str, dict[str, float | str | None]]],
|
|
) -> None:
|
|
output_path = csv_output_path(output_dir)
|
|
with output_path.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.writer(handle)
|
|
top_header = [""]
|
|
sub_header = ["Method"]
|
|
for dataset_name in datasets:
|
|
pretty_name = dataset_label(dataset_name)
|
|
top_header.extend([pretty_name] * len(CSV_SUBCOLUMNS))
|
|
sub_header.extend(CSV_SUBCOLUMNS)
|
|
|
|
writer.writerow(top_header)
|
|
writer.writerow(sub_header)
|
|
|
|
for method_name in ("RAG", "RAG + D2L", "D2L"):
|
|
row = [method_name]
|
|
for dataset_name in datasets:
|
|
method_values = summary.get(method_name, {}).get(dataset_name, {})
|
|
for subcolumn in CSV_SUBCOLUMNS:
|
|
value = method_values.get(subcolumn)
|
|
if isinstance(value, (int, float)):
|
|
row.append(f"{float(value):.4f}")
|
|
elif value is None:
|
|
row.append("")
|
|
else:
|
|
row.append(str(value))
|
|
writer.writerow(row)
|
|
|
|
print(f"Saved CSV to {output_path}")
|
|
|
|
|
|
def main() -> None:
|
|
args = build_parser().parse_args()
|
|
configure_plot_style()
|
|
|
|
allowed_datasets = set(args.datasets) if args.datasets else None
|
|
allowed_top_ks = set(args.top_ks) if args.top_ks else None
|
|
output_dir = (
|
|
args.hybrid_root / "plots" if args.output_dir is None else args.output_dir
|
|
)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
raw_runs = select_latest_runs(
|
|
args.rag_root,
|
|
want_hybrid=False,
|
|
allowed_datasets=allowed_datasets,
|
|
allowed_top_ks=allowed_top_ks,
|
|
rag_chunk_size=args.rag_chunk_size,
|
|
rag_chunk_overlap=args.rag_chunk_overlap,
|
|
rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
|
|
)
|
|
hybrid_runs = select_latest_runs(
|
|
args.hybrid_root,
|
|
want_hybrid=True,
|
|
allowed_datasets=allowed_datasets,
|
|
allowed_top_ks=allowed_top_ks,
|
|
rag_chunk_size=args.rag_chunk_size,
|
|
rag_chunk_overlap=args.rag_chunk_overlap,
|
|
rag_max_retrieved_tokens=args.rag_max_retrieved_tokens,
|
|
)
|
|
|
|
shared_keys = set(raw_runs) & set(hybrid_runs)
|
|
if not shared_keys:
|
|
raise SystemExit("Could not find any matched raw-RAG and hybrid-RAG runs.")
|
|
|
|
datasets = sorted({dataset_name for dataset_name, _ in shared_keys})
|
|
base_runs = select_latest_plain_base_runs(BASE_MODEL_EVAL_ROOT, set(datasets))
|
|
d2l_runs = select_latest_plain_d2l_runs(args.hybrid_root, set(datasets))
|
|
summary: dict[str, dict[str, dict[str, float | str | None]]] = {
|
|
"RAG": {},
|
|
"RAG + D2L": {},
|
|
"D2L": {},
|
|
}
|
|
for dataset_name in datasets:
|
|
base_run_dir = base_runs.get(dataset_name)
|
|
if base_run_dir is None:
|
|
continue
|
|
|
|
shared_top_ks = sorted(
|
|
top_k
|
|
for candidate_dataset, top_k in shared_keys
|
|
if candidate_dataset == dataset_name
|
|
)
|
|
if not shared_top_ks:
|
|
continue
|
|
|
|
raw_result_json = load_results_json(
|
|
raw_runs[(dataset_name, shared_top_ks[0])].run_dir, dataset_name
|
|
)
|
|
hybrid_result_json = load_results_json(
|
|
hybrid_runs[(dataset_name, shared_top_ks[0])].run_dir,
|
|
dataset_name,
|
|
)
|
|
metric_name = infer_metric_name(raw_result_json, dataset_name)
|
|
hybrid_metric_name = infer_metric_name(hybrid_result_json, dataset_name)
|
|
if metric_name != hybrid_metric_name:
|
|
raise ValueError(
|
|
f"Metric mismatch for {dataset_name}: raw uses {metric_name}, hybrid uses {hybrid_metric_name}."
|
|
)
|
|
base_result_json = load_results_json(base_run_dir, dataset_name)
|
|
base_metric_name = infer_metric_name(base_result_json, dataset_name)
|
|
if metric_name != base_metric_name:
|
|
raise ValueError(
|
|
f"Metric mismatch for {dataset_name}: comparison uses {metric_name}, baseline uses {base_metric_name}."
|
|
)
|
|
base_value = get_metric_value(base_run_dir, dataset_name, metric_name)
|
|
if base_value in (None, 0):
|
|
continue
|
|
|
|
d2l_run_dir = d2l_runs.get(dataset_name)
|
|
d2l_value = None
|
|
if d2l_run_dir is not None:
|
|
d2l_result_json = load_results_json(d2l_run_dir, dataset_name)
|
|
d2l_metric_name = infer_metric_name(d2l_result_json, dataset_name)
|
|
if d2l_metric_name == metric_name:
|
|
candidate_value = get_metric_value(
|
|
d2l_run_dir, dataset_name, metric_name
|
|
)
|
|
if candidate_value is not None:
|
|
d2l_value = candidate_value / base_value
|
|
|
|
paired_points: list[tuple[int, float, float]] = []
|
|
for top_k in shared_top_ks:
|
|
raw_value = get_metric_value(
|
|
raw_runs[(dataset_name, top_k)].run_dir, dataset_name, metric_name
|
|
)
|
|
hybrid_value = get_metric_value(
|
|
hybrid_runs[(dataset_name, top_k)].run_dir,
|
|
dataset_name,
|
|
metric_name,
|
|
)
|
|
if raw_value is None or hybrid_value is None:
|
|
continue
|
|
paired_points.append(
|
|
(top_k, raw_value / base_value, hybrid_value / base_value)
|
|
)
|
|
|
|
if not paired_points:
|
|
continue
|
|
|
|
summary["RAG"][dataset_name] = {
|
|
NO_CONTEXT_COLUMN: "N/A",
|
|
**{f"top-{top_k}": raw_value for top_k, raw_value, _ in paired_points},
|
|
}
|
|
summary["RAG + D2L"][dataset_name] = {
|
|
NO_CONTEXT_COLUMN: "N/A",
|
|
**{
|
|
f"top-{top_k}": hybrid_value for top_k, _, hybrid_value in paired_points
|
|
},
|
|
}
|
|
summary["D2L"][dataset_name] = {
|
|
NO_CONTEXT_COLUMN: d2l_value,
|
|
"top-1": "N/A",
|
|
"top-4": "N/A",
|
|
}
|
|
plot_confusion_matrices(
|
|
output_dir,
|
|
dataset_name,
|
|
metric_name,
|
|
paired_points,
|
|
raw_runs,
|
|
hybrid_runs,
|
|
)
|
|
|
|
write_summary_csv(output_dir, datasets, summary)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|