add functionality to sweep / cleanup

elk/plotting/command.py

@@ -22,7 +22,7 @@ class Plot:
  overwrite: bool = False
  """Whether to overwrite existing plots."""
 
- metric_type: str = None
+ metric_type: str = "auroc_estimate"
  """Name of metric to plot"""
 
  def execute(self):
@@ -38,10 +38,6 @@ def execute(self):
  else:
  sweep_paths = [root_dir / sweep for sweep in self.sweeps]
 
- if not self.metric_type:
- # ArgumentParser maps cli input --metric to metric_type
- self.metric_type = "auroc_estimate"
-
  for sweep_path in sweep_paths:
  if not sweep_path.exists():
  pretty_error(f"No sweep with name {{{sweep_path}}} found in {root_dir}")

elk/plotting/visualize.py

@@ -114,7 +114,7 @@ class TransferEvalHeatmap:
  """Class for generating heatmaps for transfer evaluation results."""
 
  layer: int
- metric_type: str = None
+ metric_type: str = ""
  ensembling: str = "full"
 
  def render(self, df: pd.DataFrame) -> go.Figure:
@@ -149,7 +149,7 @@ class TransferEvalTrend:
  evaluation."""
 
  dataset_names: list[str] | None
- metric_type: str = None
+ metric_type: str = ""
 
  def render(self, df: pd.DataFrame) -> go.Figure:
  """Render the trend plot visualization.

elk/training/sweep.py

@@ -50,6 +50,9 @@ class Sweep:
  visualize: bool = False
  """Whether to generate visualizations of the results of the sweep."""
 
+ metric_type: str = "auroc_estimate"
+ """Name of metric to plot"""
+
  name: str | None = None
 
  # A bit of a hack to add all the command line arguments from Elicit
@@ -176,4 +179,4 @@ def execute(self):
  eval.execute(highlight_color="green")
 
  if self.visualize:
- visualize_sweep(sweep_dir)
+ visualize_sweep(sweep_dir, self.metric_type)