Output eval logging (batch level)

composer/callbacks/__init__.py

@@ -9,6 +9,7 @@
 from composer.callbacks.activation_monitor import ActivationMonitor
 from composer.callbacks.checkpoint_saver import CheckpointSaver
 from composer.callbacks.early_stopper import EarlyStopper
+from composer.callbacks.eval_output_logging_callback import EvalOutputLogging
 from composer.callbacks.export_for_inference import ExportForInferenceCallback
 from composer.callbacks.free_outputs import FreeOutputs
 from composer.callbacks.generate import Generate
@@ -35,6 +36,7 @@
     'CheckpointSaver',
     'MLPerfCallback',
     'EarlyStopper',
+    'EvalOutputLogging',
     'ExportForInferenceCallback',
     'ThresholdStopper',
     'ImageVisualizer',

composer/callbacks/eval_output_logging_callback.py

@@ -0,0 +1,120 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Log model outputs and expected outputs during ICL evaluation."""
+
+import warnings
+from copy import deepcopy
+from typing import Any, Dict, List, Sequence, Union
+
+import torch
+
+from composer.core import Callback, State
+from composer.loggers import ConsoleLogger, Logger
+from composer.utils.dist import all_gather_object
+
+
+class EvalOutputLogging(Callback):
+    """Logs eval outputs for each sample of each ICL evaluation dataset.
+
+    ICL metrics are required to support caching the model's responses including information on whether model was correct.
+    Metrics are responsible for returning the results of individual datapoints in a dictionary of lists.
+    The callback will log the metric name, the depadded and detokenized input, any data stored in state.metric_outputs, and
+    any keys from the batch pased into `batch_keys_to_log`. It will do so after every eval batch.
+    """
+
+    def __init__(self, log_tokens=False, *args, **kwargs):
+        super().__init__(self, *args, **kwargs)
+        self.log_tokens = log_tokens
+        self.columns = None
+        self.name = None
+        self.rows = []
+
+    def eval_batch_end(self, state: State, logger: Logger) -> None:
+        if not isinstance(state.batch, Dict):
+            warnings.warn(f'''EvalOutputLogging only supports batches that are dictionary. \
+                Found batch for type {type(state.batch)}. \
+                Not logging eval outputs.''')
+            return
+
+        assert state.outputs is not None
+        assert state.metric_outputs is not None
+        logging_dict: Dict[str, Union[List[Any], torch.Tensor, Sequence[torch.Tensor]]] = deepcopy(state.metric_outputs)
+
+        # If batch mode is not generate, outputs will be logits
+        if state.batch['mode'] == 'generate':
+            # Outputs are already detokenized
+            logging_dict['outputs'] = state.outputs
+
+        input_ids = state.batch['input_ids']
+        logged_input = []
+        assert state.dataloader is not None
+
+        # Depad and decode input_ids
+        for input_list in input_ids.tolist():
+            depadded_input = [
+                tok for tok in input_list
+                if tok != state.dataloader.dataset.pad_tok_id  # pyright: ignore[reportGeneralTypeIssues]
+            ]
+            logged_input.append(
+                state.dataloader.dataset.tokenizer.decode(depadded_input))  # pyright: ignore[reportGeneralTypeIssues]
+        logging_dict['input'] = logged_input
+
+        # Log token indices if toggled
+        if self.log_tokens:
+            logging_dict['input_tokens'] = input_ids.tolist()
+            if not state.batch['mode'] == 'generate':
+                if isinstance(state.outputs, torch.Tensor):  # pyright
+                    logging_dict['label_tokens'] = state.outputs.tolist()
+
+        # Add run_name as a column
+        run_name_list = [state.run_name for _ in range(0, len(logging_dict['input']))]
+        logging_dict['run_name'] = run_name_list
+
+        # NOTE: This assumes _any_ tensor logged are tokens to be decoded.
+        #       This might not be true if, for example, logits are logged.
+
+        # Detokenize data in rows
+        for key, value in logging_dict.items():
+            # All types in list are the same
+            if isinstance(value[0], torch.Tensor):
+                logging_dict[key] = [
+                    state.dataloader.dataset.tokenizer.decode(t)  # pyright: ignore[reportGeneralTypeIssues]
+                    for t in value
+                ]
+            elif isinstance(value[0], list):
+                if isinstance(value[0][0], torch.Tensor):
+                    logging_dict[key] = [
+                        [
+                            state.dataloader.dataset.tokenizer.decode(  # pyright: ignore[reportGeneralTypeIssues]
+                                choice) for choice in t
+                        ] for t in value
+                    ]
+
+        # Convert logging_dict from kv pairs of column name and column values to a list of rows
+        # Example:
+        # logging_dict = {"a": ["1a", "2a"], "b": ["1b", "2b"]}
+        # will become
+        # columns = {"a", "b"}, rows = [["1a", "1b"], ["2a", "2b"]]
+        columns = list(logging_dict.keys())
+        rows = [list(item) for item in zip(*logging_dict.values())]
+
+        assert state.dataloader_label is not None
+        if not self.name:
+            # If only running eval, step will be 0
+            # If running training, step will be current training step
+            step = state.timestamp.batch.value
+            self.name = f'{state.dataloader_label}_step_{step}'
+            self.columns = columns
+        self.rows.extend(rows)
+
+    def eval_end(self, state: State, logger: Logger) -> None:
+        list_of_rows = all_gather_object(self.rows)
+        rows = [row for rows in list_of_rows for row in rows]
+        for dest_logger in logger.destinations:
+            if not isinstance(dest_logger, ConsoleLogger):
+                dest_logger.log_table(self.columns, rows, name=self.name, step=state.timestamp.batch.value)
+
+        self.rows = []
+        self.name = None
+        self.columns = None

composer/core/state.py

@@ -518,6 +518,8 @@ def __init__(
         self.eval_metric_values: Dict[str, float] = {}
         self.total_loss_dict: Dict[str, float] = {}
 
+        self.metric_outputs: Dict[str, Any] = {}
+
     def _dataset_of(self, dataloader: Optional[Union[Evaluator, DataSpec, DataLoader, Iterable]]) -> Optional[Dataset]:
         """Get the dataset contained by the given dataloader-like object.
 

composer/loggers/in_memory_logger.py

@@ -83,8 +83,10 @@ def log_table(self,
             raise MissingConditionalImportError(extra_deps_group='pandas',
                                                 conda_package='pandas',
                                                 conda_channel='conda-forge') from e
-        table = pd.DataFrame.from_records(data=rows, columns=columns).to_json(orient='split', index=False)
-        assert isinstance(table, str)
+        table = pd.DataFrame.from_records(data=rows, columns=columns).to_json(orient='split',
+                                                                              index=False,
+                                                                              force_ascii=False)
+        assert table is not None
         self.tables[name] = table
 
     def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:

composer/loggers/wandb_logger.py

@@ -109,6 +109,8 @@ def __init__(
         self.run_dir: Optional[str] = None
         self.run_url: Optional[str] = None
 
+        self.table_dict = {}
+
     def _set_is_in_atexit(self):
         self._is_in_atexit = True
 
@@ -125,7 +127,7 @@ def log_table(self,
         if self._enabled:
             import wandb
             table = wandb.Table(columns=columns, rows=rows)
-            wandb.log({name: table}, step)
+            wandb.log({name: table}, step=step)
 
     def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
         if self._enabled: