resolved circular import

elk/evaluation/evaluate.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal
 
 import pandas as pd
 import torch
@@ -9,6 +10,7 @@
 from ..files import elk_reporter_dir
 from ..metrics import evaluate_preds
 from ..run import Run
+from ..training.multi_reporter import AnyReporter, MultiReporter
 from ..utils import Color
 
 
@@ -38,39 +40,69 @@
 
  experiment_dir = elk_reporter_dir() / self.source
 
- reporter_path = experiment_dir / "reporters" / f"layer_{layer}.pt"
- reporter = torch.load(reporter_path, map_location=device)
+ def load_reporter() -> AnyReporter | MultiReporter:
+ # check if experiment_dir / "reporters" has .pt files
+ first = next((experiment_dir / "reporters").iterdir())
+ if not first.suffix == ".pt":
+ return MultiReporter.load(
+ experiment_dir / "reporters", layer, device=device
+ )
+ else:
+ path = experiment_dir / "reporters" / f"layer_{layer}.pt"
+ return torch.load(path, map_location=device)
+
+ reporter = load_reporter()
 
  row_bufs = defaultdict(list)
- for ds_name, (val_h, val_gt, _) in val_output.items():
- meta = {"dataset": ds_name, "layer": layer}
-
- val_credences = reporter(val_h)
- for mode in ("none", "partial", "full"):
- row_bufs["eval"].append(
- {
- **meta,
- "ensembling": mode,
- **evaluate_preds(val_gt, val_credences, mode).to_dict(),
- }
- )
 
- lr_dir = experiment_dir / "lr_models"
- if not self.skip_supervised and lr_dir.exists():
- with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
- lr_models = torch.load(f, map_location=device)
- if not isinstance(lr_models, list): # backward compatibility
- lr_models = [lr_models]
-
- for i, model in enumerate(lr_models):
- model.eval()
- row_bufs["lr_eval"].append(
- {
- "ensembling": mode,
- "inlp_iter": i,
- **meta,
- **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
- }
- )
+ def eval_all(
+ reporter: AnyReporter | MultiReporter,
+ prompt_index: int | Literal["multi"] | None = None,
+ ):
+ prompt_index = (
+ {"prompt_index": prompt_index} if prompt_index is not None else {}
+ )
+ for ds_name, (val_h, val_gt, _) in val_output.items():
+ meta = {"dataset": ds_name, "layer": layer}
+
+ val_credences = reporter(val_h)
+ for mode in ("none", "partial", "full"):
+ row_bufs["eval"].append(
+ {
+ **meta,
+ "ensembling": mode,
+ **evaluate_preds(val_gt, val_credences, mode).to_dict(),
+ **prompt_index,
+ }
+ )
+
+ lr_dir = experiment_dir / "lr_models"
+ if not self.skip_supervised and lr_dir.exists():
+ with open(lr_dir / f"layer_{layer}.pt", "rb") as f:
+ lr_models = torch.load(f, map_location=device)
+ if not isinstance(
+ lr_models, list
+ ): # backward compatibility
+ lr_models = [lr_models]
+
+ for i, model in enumerate(lr_models):
+ model.eval()
+ row_bufs["lr_eval"].append(
+ {
+ "ensembling": mode,
+ "inlp_iter": i,
+ **meta,
+ **evaluate_preds(
+ val_gt, model(val_h), mode
+ ).to_dict(),
+ }
+ )
+
+ if isinstance(reporter, MultiReporter):
+ for prompt_index, single_reporter in enumerate(reporter.reporters):
+ eval_all(single_reporter, prompt_index)
+ eval_all(reporter, "multi")
+ else:
+ eval_all(reporter)
 
  return {k: pd.DataFrame(v) for k, v in row_bufs.items()}

elk/training/multi_reporter.py

@@ -0,0 +1,44 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch as t
+
+from elk.training import CcsReporter
+from elk.training.common import Reporter
+
+AnyReporter = CcsReporter | Reporter
+
+
+@dataclass
+class ReporterTrainResult:
+ reporter: AnyReporter
+ train_loss: float | None
+
+
+class MultiReporter:
+ def __init__(self, reporter_results: list[ReporterTrainResult]):
+ self.reporter_results: list[ReporterTrainResult] = reporter_results
+ self.reporters = [r.reporter for r in reporter_results]
+ train_losses = (
+ [r.train_loss for r in reporter_results]
+ if reporter_results[0].train_loss is not None
+ else None
+ )
+ self.train_loss = (
+ sum(train_losses) / len(train_losses) if train_losses is not None else None
+ )
+
+ def __call__(self, h):
+ credences = [r(h) for r in self.reporters]
+ return t.stack(credences).mean(dim=0)
+
+ @staticmethod
+ def load(path: Path, layer: int, device: str):
+ prompt_folders = [p for p in path.iterdir() if p.is_dir()]
+ reporters = []
+ for folder in prompt_folders:
+ path = folder / "reporters" / f"layer_{layer}.pt"
+ reporter = t.load(path, map_location=device)
+ reporters.append(reporter)
+ # TODO for now I don't care about the train losses
+ return MultiReporter([ReporterTrainResult(r, None) for r in reporters])

elk/training/train.py

@@ -17,34 +17,9 @@
 from ..training.supervised import train_supervised
 from . import Classifier
 from .ccs_reporter import CcsConfig, CcsReporter
-from .common import FitterConfig, Reporter
+from .common import FitterConfig
 from .eigen_reporter import EigenFitter, EigenFitterConfig
-
-AnyReporter = CcsReporter | Reporter
-
-
-@dataclass
-class ReporterTrainResult:
- reporter: AnyReporter
- train_loss: float | None
-
-
-class MultiReporter:
- def __init__(self, reporter_results: list[ReporterTrainResult]):
- self.reporter_results: list[ReporterTrainResult] = reporter_results
- self.reporters = [r.reporter for r in reporter_results]
- train_losses = (
- [r.train_loss for r in reporter_results]
- if reporter_results[0].train_loss is not None
- else None
- )
- self.train_loss = (
- sum(train_losses) / len(train_losses) if train_losses is not None else None
- )
-
- def __call__(self, h):
- credences = [r(h) for r in self.reporters]
- return torch.stack(credences).mean(dim=0)
+from .multi_reporter import AnyReporter, MultiReporter, ReporterTrainResult
 
 
 def evaluate_and_save(
@@ -67,7 +42,7 @@
  ):
  val_credences = reporter(val_h)
  train_credences = reporter(train_h)
  prompt_index = {"prompt_index": prompt_index}
  for mode in ("none", "partial", "full"):
  row_bufs["eval"].append(
  {
@@ -75,7 +50,7 @@
  "ensembling": mode,
  **evaluate_preds(val_gt, val_credences, mode).to_dict(),
  "train_loss": train_loss,
  **prompt_index,
  }
  )
 
@@ -85,7 +60,7 @@
  "ensembling": mode,
  **evaluate_preds(train_gt, train_credences, mode).to_dict(),
  "train_loss": train_loss,
  **prompt_index,
  }
  )
 
@@ -95,7 +70,7 @@
  **meta,
  "ensembling": mode,
  **evaluate_preds(val_gt, val_lm_preds, mode).to_dict(),
  **prompt_index,
  }
  )
 
@@ -105,7 +80,7 @@
  **meta,
  "ensembling": mode,
  **evaluate_preds(train_gt, train_lm_preds, mode).to_dict(),
  **prompt_index,
  }
  )
 
@@ -116,7 +91,7 @@
  "ensembling": mode,
  "inlp_iter": i,
  **evaluate_preds(val_gt, model(val_h), mode).to_dict(),
  **prompt_index,
  }
  )
 
@@ -322,5 +297,5 @@
  )
 
  return evaluate_and_save(
  train_loss, maybe_multi_reporter, train_dict, val_dict, lr_models, layer
  )