Select metric choice for visualization

README.md

@@ -57,11 +57,11 @@ together. You can also add a `--visualize` flag to visualize the results of the
 elk sweep --models gpt2-{medium,large,xl} --datasets imdb amazon_polarity --add_pooled
 ```
 
-If you just do `elk plot`, it will plot the results from the most recent sweep.
-If you want to plot a specific sweep, you can do so with:
+If you just do `elk plot`, it will plot the results of AUROC from the most recent sweep.
+If you want to plot a specific sweep, with a specific metric type, you can do so with:
 
 ```bash
-elk plot {sweep_name}
+elk plot {sweep_name} --metric acc_estimate
 ```
 
 ## Caching

elk/plotting/command.py

@@ -22,6 +22,9 @@ class Plot:
  overwrite: bool = False
  """Whether to overwrite existing plots."""
 
+ metric_type: str = "auroc_estimate"
+ """Name of metric to plot"""
+
  def execute(self):
  root_dir = sweeps_dir()
 
@@ -47,4 +50,4 @@ def execute(self):
  if self.overwrite:
  shutil.rmtree(sweep_path / "viz")
 
- visualize_sweep(sweep_path)
+ visualize_sweep(sweep_path, self.metric_type)

elk/plotting/visualize.py

@@ -47,7 +47,7 @@ def render(
  shared_yaxes=True,
  vertical_spacing=0.1,
  x_title="Layer",
- y_title="AUROC",
+ y_title=f"{sweep.metric_type}",
  )
  color_map = dict(zip(ensembles, qualitative.Plotly))
 
@@ -56,7 +56,7 @@ def render(
  if with_transfer: # TODO write tests
  ensemble_data = ensemble_data.groupby(
  ["eval_dataset", "layer", "ensembling"], as_index=False
- ).agg({"auroc_estimate": "mean"})
+ ).agg({f"{sweep.metric_type}": "mean"})
  else:
  ensemble_data = ensemble_data[
  ensemble_data["eval_dataset"] == ensemble_data["train_dataset"]
@@ -75,7 +75,7 @@ def render(
  fig.add_trace(
  go.Scatter(
  x=dataset_data["layer"],
- y=dataset_data["auroc_estimate"],
+ y=dataset_data[f"{sweep.metric_type}"],
  mode="lines",
  name=ensemble,
  showlegend=False
@@ -95,7 +95,7 @@ def render(
  legend=dict(
  title="Ensembling",
  ),
- title=f"AUROC Trend: {self.model_name}",
+ title=f"{sweep.metric_type} Trend: {self.model_name}",
  )
  if write:
  fig.write_image(
@@ -114,7 +114,7 @@ class TransferEvalHeatmap:
  """Class for generating heatmaps for transfer evaluation results."""
 
  layer: int
- score_type: str = "auroc_estimate"
+ metric_type: str = ""
  ensembling: str = "full"
 
  def render(self, df: pd.DataFrame) -> go.Figure:
@@ -129,27 +129,28 @@ def render(self, df: pd.DataFrame) -> go.Figure:
  model_name = df["eval_dataset"].iloc[0] # infer model name
  # TODO: validate
  pivot = pd.pivot_table(
- df, values=self.score_type, index="eval_dataset", columns="train_dataset"
+ df, values=self.metric_type, index="eval_dataset", columns="train_dataset"
  )
 
  fig = px.imshow(pivot, color_continuous_scale="Viridis", text_auto=True)
 
  fig.update_layout(
  xaxis_title="Train Dataset",
  yaxis_title="Transfer Dataset",
- title=f"AUROC Score Heatmap: {model_name} | Layer {self.layer}",
+ title=f"{self.metric_type} Score Heatmap: {model_name} \
+ | Layer {self.layer}",
  )
 
  return fig
 
 
 @dataclass
 class TransferEvalTrend:
- """Class for generating line plots for the trend of AUROC scores in transfer
+ """Class for generating line plots for the trend of metric scores in transfer
  evaluation."""
 
  dataset_names: list[str] | None
- score_type: str = "auroc_estimate"
+ metric_type: str = ""
 
  def render(self, df: pd.DataFrame) -> go.Figure:
  """Render the trend plot visualization.
@@ -164,14 +165,14 @@ def render(self, df: pd.DataFrame) -> go.Figure:
  if self.dataset_names is not None:
  df = self._filter_transfer_datasets(df, self.dataset_names)
  pivot = pd.pivot_table(
- df, values=self.score_type, index="layer", columns="eval_dataset"
+ df, values=self.metric_type, index="layer", columns="eval_dataset"
  )
 
  fig = px.line(pivot, color_discrete_sequence=px.colors.qualitative.Plotly)
  fig.update_layout(
  xaxis_title="Layer",
- yaxis_title="AUROC Score",
- title=f"AUROC Score Trend: {model_name}",
+ yaxis_title=f"{self.metric_type} Score",
+ title=f"{self.metric_type} Score Trend: {model_name}",
  )
 
  avg = pivot.mean(axis=1)
@@ -244,17 +245,16 @@ def render_and_save(
  self,
  sweep: "SweepVisualization",
  dataset_names: list[str] | None = None,
- score_type="auroc_estimate",
  ensembling="full",
  ) -> None:
  """Render and save the visualization for the model.
 
  Args:
  sweep: The SweepVisualization instance.
  dataset_names: List of dataset names to include in the visualization.
- score_type: The type of score to display.
  ensembling: The ensembling option to consider.
  """
+ metric_type = sweep.metric_type
  df = self.df
  model_name = self.model_name
  layer_min, layer_max = df["layer"].min(), df["layer"].max()
@@ -264,10 +264,10 @@ def render_and_save(
  for layer in range(layer_min, layer_max + 1):
  filtered = df[(df["layer"] == layer) & (df["ensembling"] == ensembling)]
  fig = TransferEvalHeatmap(
- layer, score_type=score_type, ensembling=ensembling
+ layer, metric_type=metric_type, ensembling=ensembling
  ).render(filtered)
  fig.write_image(file=model_path / f"{layer}.png")
- fig = TransferEvalTrend(dataset_names).render(df)
+ fig = TransferEvalTrend(dataset_names, metric_type=metric_type).render(df)
  fig.write_image(file=model_path / "transfer_eval_trend.png")
 
  @staticmethod
@@ -288,6 +288,7 @@ class SweepVisualization:
  path: Path
  datasets: list[str]
  models: dict[str, ModelVisualization]
+ metric_type: str
 
  def model_names(self) -> list[str]:
  """Get the names of all models in the sweep.
@@ -323,7 +324,7 @@ def _get_model_paths(sweep_path: Path) -> list[Path]:
  return folders
 
  @classmethod
- def collect(cls, sweep_path: Path) -> "SweepVisualization":
+ def collect(cls, sweep_path: Path, metric_type: str) -> "SweepVisualization":
  """Collect the evaluation data for a sweep.
 
  Args:
@@ -348,7 +349,9 @@ def collect(cls, sweep_path: Path) -> "SweepVisualization":
  }
  df = pd.concat([model.df for model in models.values()], ignore_index=True)
  datasets = list(df["eval_dataset"].unique())
- return cls(sweep_name, df, sweep_viz_path, datasets, models)
+ return cls(
+ sweep_name, df, sweep_viz_path, datasets, models, metric_type=metric_type
+ )
 
  def render_and_save(self):
  """Render and save all visualizations for the sweep."""
@@ -368,14 +371,11 @@ def render_multiplots(self, write=False):
  for model in self.models
  ]
 
- def render_table(
- self, score_type="auroc_estimate", display=True, write=False
- ) -> pd.DataFrame:
+ def render_table(self, display=True, write=False) -> pd.DataFrame:
  """Render and optionally write the score table.
 
  Args:
  layer: The layer number (from last layer) to include in the score table.
- score_type: The type of score to include in the table.
  display: Flag indicating whether to display the table to stdout.
  write: Flag indicating whether to write the table to a file.
 
@@ -387,15 +387,15 @@ def render_table(
  # For each model, we use the layer whose mean AUROC is the highest
  best_layers, model_dfs = [], []
  for _, model_df in df.groupby("model_name"):
- best_layer = model_df.groupby("layer").auroc_estimate.mean().argmax()
+ best_layer = model_df.groupby("layer")[self.metric_type].mean().argmax()
 
  best_layers.append(best_layer)
  model_dfs.append(model_df[model_df["layer"] == best_layer])
 
  pivot_table = pd.concat(model_dfs).pivot_table(
  index="eval_dataset",
  columns="model_name",
- values=score_type,
+ values=self.metric_type,
  margins=True,
  margins_name="Mean",
  )
@@ -416,14 +416,14 @@ def render_table(
  console.print(table)
 
  if write:
- pivot_table.to_csv(f"score_table_{score_type}.csv")
+ pivot_table.to_csv(f"score_table_{self.metric_type}.csv")
  return pivot_table
 
 
-def visualize_sweep(sweep_path: Path):
+def visualize_sweep(sweep_path: Path, metric_type: str):
  """Visualize a sweep by generating and saving the visualizations.
 
  Args:
  sweep_path: The path to the sweep data directory.
  """
- SweepVisualization.collect(sweep_path).render_and_save()
+ SweepVisualization.collect(sweep_path, metric_type).render_and_save()

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)