From 2d6f08397eaa9cd7e3f22cdb9a90dd7f46c3ed20 Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Tue, 17 Jun 2025 11:35:04 +0000
Subject: [PATCH] llm comparator

---
 src/ctx_to_lora/eval_utils.py | 11 ++++++-----
 src/ctx_to_lora/viz_utils.py  | 23 +++++++++++++++++++----
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py
index 3e31272..10d5b26 100644
--- a/src/ctx_to_lora/eval_utils.py
+++ b/src/ctx_to_lora/eval_utils.py
@@ -599,10 +599,10 @@ def eval_teacher_forcing(
     for ds_name, ds in datasets.items():
         split_name = f"{split}_{ds_name}" if ds_name else split
         if "/" in split_name:
-            os.makedirs(
-                f"{eval_trainer.args.output_dir}/{split_name.split('/')[0]}",
-                exist_ok=True,
+            split_dir = os.path.join(
+                eval_trainer.args.output_dir, os.path.dirname(split_name)
             )
+            os.makedirs(split_dir, exist_ok=True)
         if remove_context:
             split_name += "_no_context"
 
@@ -735,7 +735,8 @@ def evaluate(
         for ds_name, ds in datasets.items():
             val_indices = np.random.permutation(len(ds))[:max_eval_samples_per_ds]
             datasets[ds_name] = ds.select(val_indices)
-            answers[ds_name] = answers[ds_name].select(val_indices)
+            if ds_name in answers:
+                answers[ds_name] = answers[ds_name].select(val_indices)
 
         print(f"Datasets: {datasets}")
         print(f"Answers: {answers}")
@@ -879,7 +880,7 @@ def run_eval(
         else:
             args.test_ds_names = datasets
 
-    evaluate(
+    return evaluate(
         checkpoint_path,
         model_name_or_path,
         eval_batch_size,
diff --git a/src/ctx_to_lora/viz_utils.py b/src/ctx_to_lora/viz_utils.py
index af88286..1e503db 100644
--- a/src/ctx_to_lora/viz_utils.py
+++ b/src/ctx_to_lora/viz_utils.py
@@ -28,11 +28,18 @@ def convert_to_llm_comparator_format(
         files_model = sorted(
             glob(f"{eval_model_dir}/**/*no_context*generated*.jsonl", recursive=True)
         )
+        files_model = {
+            file.split("/")[-1].replace("_no_context", ""): file for file in files_model
+        }
     else:
         files_model = sorted(
-            glob(f"{eval_model_dir}/**/*[!no_context]*generated*.jsonl", recursive=True)
+            glob(f"{eval_model_dir}/**/*generated*.jsonl", recursive=True)
         )
-    files_model = {file.split("/")[-1]: file for file in files_model}
+        files_model = {
+            file.split("/")[-1]: file
+            for file in files_model
+            if "no_context" not in file
+        }
 
     overlap_files = set(files_checkpoint) & set(files_model)
     if not overlap_files:
@@ -301,11 +308,19 @@ if __name__ == "__main__":
     # Plot and save confusion matrix
     checkpoint_name = "/".join(eval_checkpoint_dir.strip("/").split("/")[-2:])
     plot_title = f"Model Comparison: {checkpoint_name} vs {model_dir}"
-    save_path = f"{eval_checkpoint_dir}/confusion_matrix.pdf"
+    fname = (
+        "confusion_matrix_no_context" if compare_to_no_context else "confusion_matrix"
+    )
+    save_path = f"{eval_checkpoint_dir}/{fname}.pdf"
     # TODO: add resilience + boost rate
     fig = plot_confusion_matrix(confusion_matrix, title=plot_title, save_path=save_path)
     # fig.show()
 
     # Save the comparator JSON to a file
-    with open(f"{eval_checkpoint_dir}/comparator.json", "w") as f:
+    comparator_fname = (
+        "comparator_no_context.json" if compare_to_no_context else "comparator.json"
+    )
+    with open(f"{eval_checkpoint_dir}/{comparator_fname}", "w") as f:
         json.dump(comparator_json, f)
+    print(f"Comparator JSON saved to: {eval_checkpoint_dir}/{comparator_fname}")
+    print("Done!")