fix: per-module aggregation in result analysis with several hookpoints

delphi/log/result_analysis.py

@@ -193,7 +193,11 @@ def parse_score_file(path: Path) -> pd.DataFrame:
         if not score_type_dir.is_dir():
             continue
         for module in modules:
-            for file in score_type_dir.glob(f"*{module}*"):
+            # Score files are named f"{module}_latent{idx}" (see Latent.__str__);
+            # match the full module name so that e.g. "layers.1" does not also
+            # pick up the files of "layers.10"
+            safe_module = module.replace("/", "--")
+            for file in score_type_dir.glob(f"{safe_module}_latent*"):
                 latent_idx = int(file.stem.split("latent")[-1])
 
                 latent_df = parse_score_file(file)
@@ -203,7 +207,7 @@ def parse_score_file(path: Path) -> pd.DataFrame:
                 if counts:
                     latent_df["firing_count"] = (
                         counts[module][latent_idx].item()
-                        if latent_idx in counts[module]
+                        if latent_idx < len(counts[module])
                         else None
                     )
 
@@ -234,7 +238,10 @@ def frequency_weighted_f1(
     for module in latent_df["module"].unique():
         module_df = latent_df[latent_df["module"] == module]
 
-        firing_weights = counts[module][module_df["latent_idx"]].float()
+        # Index with the values: a pandas Series would be indexed by its labels,
+        # which do not start at 0 for any module but the first
+        latent_indices = torch.as_tensor(module_df["latent_idx"].to_numpy())
+        firing_weights = counts[module][latent_indices].float()
         total_weight = firing_weights.sum()
         if total_weight == 0:
             continue

tests/test_log/test_result_analysis.py

@@ -1,7 +1,10 @@
+import json
+
 import numpy as np
 import pandas as pd
+import torch
 
-from delphi.log.result_analysis import compute_confusion
+from delphi.log.result_analysis import compute_confusion, get_agg_metrics, load_data
 
 
 def test_failed_count_counts_dropped_predictions():
@@ -46,3 +49,54 @@ def test_failed_count_zero_when_all_predictions_valid():
 
     assert conf["total_examples"] == 2
     assert conf["failed_count"] == 0
+
+
+def _example(activating: bool, prediction: bool) -> dict:
+    return {
+        "str_tokens": ["a"],
+        "distance": 0,
+        "activating": activating,
+        "prediction": prediction,
+        "probability": None,
+        "correct": activating == prediction,
+        "activations": [1.0],
+    }
+
+
+def _write_scores(tmp_path):
+    """Detection scores for latent 0 of "layers.1" (all correct) and of
+    "layers.10" (all wrong), whose name starts with "layers.1"."""
+    scores = tmp_path / "scores" / "detection"
+    scores.mkdir(parents=True)
+    (tmp_path / "log").mkdir()
+    (scores / "layers.1_latent0.txt").write_text(
+        json.dumps([_example(True, True), _example(False, False)])
+    )
+    (scores / "layers.10_latent0.txt").write_text(
+        json.dumps([_example(True, False), _example(False, True)])
+    )
+    counts = {"layers.1": torch.tensor([5, 3]), "layers.10": torch.tensor([7, 2])}
+    torch.save(counts, tmp_path / "log" / "hookpoint_firing_counts.pt")
+    return tmp_path / "scores"
+
+
+def test_load_data_keeps_modules_with_shared_prefix_apart(tmp_path):
+    latent_df, _ = load_data(_write_scores(tmp_path), ["layers.1", "layers.10"])
+
+    assert latent_df.groupby("module").size().to_dict() == {
+        "layers.1": 2,
+        "layers.10": 2,
+    }
+    layers_1 = latent_df[latent_df["module"] == "layers.1"]
+    assert (layers_1["activating"] == layers_1["prediction"]).all()
+    # firing counts are looked up by latent index
+    assert layers_1["firing_count"].tolist() == [5, 5]
+
+
+def test_frequency_weighted_f1_with_several_modules(tmp_path):
+    latent_df, counts = load_data(_write_scores(tmp_path), ["layers.1", "layers.10"])
+
+    metrics = get_agg_metrics(latent_df, counts)
+
+    # per-module F1 is 1.0 for "layers.1" and 0.0 for "layers.10"
+    assert metrics["weighted_f1"].tolist() == [0.5]