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,81 @@
+# 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 torch.utils.data import DataLoader, Dataset
+
+
+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, *args, **kwargs):
+        super().__init__(self, *args, **kwargs)
+        self.warn_batch_is_not_dict = True
+
+    def eval_batch_end(self, state: State, logger: Logger) -> None:
+        if not isinstance(state.batch, Dict):
+            if self.warn_batch_is_not_dict:
+                warnings.warn(f'''EvalOutputLogging only supports batchs that are dictionary. \
+                    Found batch for type {type(state.batch)}. \
+                    Not logging eval outputs.''')
+                self.warn_batch_is_not_dict = False
+            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 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
+
+        # Get column names
+        columns = list(logging_dict.keys())
+        # Convert logging_dict from kv pairs of column name and column values to a list of rows
+        rows = [list(item) for item in zip(*logging_dict.values())]
+
+        # 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
+        rows = [
+            [
+                state.dataloader.dataset.tokenizer.decode(x)  # pyright: ignore[reportGeneralTypeIssues]
+                if isinstance(x, torch.Tensor) else x for x in row
+            ] for row in rows
+        ]
+
+        assert state.dataloader_label is not None
+        step = state.timestamp.batch.value
+        name = f'{state.dataloader_label}_step_{step}'
+        for dest_logger in logger.destinations:
+            if not isinstance(dest_logger, ConsoleLogger):
+                dest_logger.log_table(columns, rows, name=name, step=state.timestamp.batch.value)

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
 
@@ -124,8 +126,15 @@ def log_table(self,
                   step: Optional[int] = None) -> None:
         if self._enabled:
             import wandb
-            table = wandb.Table(columns=columns, rows=rows)
-            wandb.log({name: table}, step)
+            if name in self.table_dict:
+                for row in rows:
+                    self.table_dict[name].add_data(*row)
+            else:
+                table = wandb.Table(columns=columns, rows=rows)
+                self.table_dict[name] = table
+            # Need to do this copy because apparently wandb table logging is broken LOL
+            # https://github.com/wandb/wandb/issues/2981#issuecomment-1458447291
+            wandb.log({name: copy.copy(self.table_dict[name])}, step=step)
 
     def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
         if self._enabled:

composer/metrics/nlp.py

@@ -3,12 +3,14 @@
 
 """A collection of common torchmetrics for NLP tasks."""
 
+import copy
+import functools
 import logging
 import os
 import re
 import string
 import warnings
-from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -201,6 +203,31 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.needs_batch = True
 
+    def _wrap_update(self, update: Callable) -> Callable:
+
+        @functools.wraps(update)
+        def wrapped_func(*args: Any, **kwargs: Any) -> None:
+            self._computed = None
+            self._update_count += 1
+            with torch.set_grad_enabled(self._enable_grad):
+                try:
+                    update_result = update(*args, **kwargs)
+                except RuntimeError as err:
+                    if 'Expected all tensors to be on' in str(err):
+                        raise RuntimeError(
+                            'Encountered different devices in metric calculation (see stacktrace for details).'
+                            ' This could be due to the metric class not being on the same device as input.'
+                            f' Instead of `metric={self.__class__.__name__}(...)` try to do'
+                            f' `metric={self.__class__.__name__}(...).to(device)` where'
+                            ' device corresponds to the device of the input.') from err
+                    raise err
+
+            if self.compute_on_cpu:
+                self._move_list_states_to_cpu()
+            return update_result
+
+        return wrapped_func
+
     def update(self,
                batch: dict,
                output_logits: Optional[torch.Tensor] = None,
@@ -274,6 +301,12 @@ def __init__(self, dist_sync_on_step: bool = False):
         super().__init__(dist_sync_on_step=dist_sync_on_step)
         self.add_state('correct', default=torch.tensor(0.), dist_reduce_fx='sum')
         self.add_state('total', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.metric_result_dict = {
+            'cleaned_output': [],
+            'original_label': [],
+            'cleaned_label': [],
+            'result': [],
+        }
 
     def normalize_answer(self, answer: str):
         """Lower text and remove punctuation, articles and extra whitespace.
@@ -301,8 +334,10 @@ def replace_underscore(text: str) -> str:
 
     def update(self, outputs: List[str], labels: List[List[str]], batch: Dict[str, Any]):
         cot_delimiter = batch.get('cot_delimiter', '')
+
         do_normalization = batch.get('do_normalization', True)
         stopping_criteria = batch.get('stopping_criteria', None)
+        metric_result_dict = copy.deepcopy(self.metric_result_dict)
         for sample_output, sample_labels in zip(outputs, labels):
             final_answer = sample_output
 
@@ -319,10 +354,20 @@ def update(self, outputs: List[str], labels: List[List[str]], batch: Dict[str, A
                 cleaned_final_answer = final_answer
                 cleaned_sample_labels = set(sample_labels)
 
+            metric_result_dict['original_label'].append(sample_labels)
+            metric_result_dict['cleaned_output'].append(cleaned_final_answer)
+            metric_result_dict['cleaned_label'].append(cleaned_sample_labels)
+
             if any(cleaned_final_answer.startswith(label) for label in cleaned_sample_labels):
                 self.correct += torch.tensor(1.0)
+                metric_result_dict['result'].append(1)
+            else:
+                metric_result_dict['result'].append(0)
+
             self.total += torch.tensor(1.0)
 
+        return metric_result_dict
+
     def compute(self):
         assert isinstance(self.correct, Tensor)
         assert isinstance(self.total, Tensor)
@@ -358,6 +403,7 @@ def __init__(self, dist_sync_on_step: bool = False):
         super().__init__(dist_sync_on_step=dist_sync_on_step)
         self.add_state('correct', default=torch.tensor(0.), dist_reduce_fx='sum')
         self.add_state('total', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.metric_result_dict = {'context': [], 'label': [], 'output': [], 'result': []}
 
     def update(self,
                batch: dict,
@@ -369,13 +415,23 @@ def update(self,
                                                                      labels=labels,
                                                                      outputs=outputs)
 
+        metric_result_dict = copy.deepcopy(self.metric_result_dict)
         for batch_idx, cont_idx in enumerate(batch['continuation_indices']):
             cont_tok_pred = outputs[batch_idx].index_select(dim=0, index=cont_idx - 1).argmax(dim=-1)
             cont_tok_targ = labels[batch_idx].index_select(dim=0, index=cont_idx - 1)
+            # TODO: okay to do context_tok here? or do we wanna do that in the logger?
+            metric_result_dict['context'].append(batch['input_ids'][batch_idx][:cont_idx[0]])
+            metric_result_dict['label'].append(cont_tok_targ)
+            metric_result_dict['output'].append(cont_tok_pred)
+
+            correct = (cont_tok_pred == cont_tok_targ).all().int()
+            self.correct += correct
+            metric_result_dict['result'].append(int(correct))
 
-            self.correct += (cont_tok_pred == cont_tok_targ).all().int()
             self.total += torch.tensor(1.0)
 
+        return metric_result_dict
+
     def compute(self):
         assert isinstance(self.correct, Tensor)
         assert isinstance(self.total, Tensor)
@@ -409,6 +465,7 @@ def __init__(self, dist_sync_on_step: bool = False):
         super().__init__(dist_sync_on_step=dist_sync_on_step)
         self.add_state('correct', default=torch.tensor(0.0), dist_reduce_fx='sum')
         self.add_state('total', default=torch.tensor(0.0), dist_reduce_fx='sum')
+        self.metric_result_dict = {'context': [], 'correct_choice': [], 'selected_choice': [], 'result': []}
 
     def update(self,
                batch: dict,
@@ -430,14 +487,31 @@ def update(self,
             perplexity = torch.exp(cross_entropy)
             perplexities.append(perplexity)
 
+        metric_result_dict = copy.deepcopy(self.metric_result_dict)
         for (start, end), gold_idx in zip(batch['choice_groupings'], batch['gold_indices']):
             subset = perplexities[start:end]
             idx_min = subset.index(min(subset))
-
             if idx_min == gold_idx:
                 self.correct += torch.tensor(1.0)
+                metric_result_dict['result'].append(1)
+            else:
+                metric_result_dict['result'].append(0)
+
+            question = batch['input_ids'][start][:batch['continuation_indices'][start][0]]
+
+            # TODO: Seems broken for schema tasks
+            correct_choice = batch['input_ids'][start:end][gold_idx][batch['continuation_indices'][start:end][gold_idx][
+                0]:batch['continuation_indices'][start:end][gold_idx][-1] + 1]
+            selected_choice = batch['input_ids'][start:end][idx_min][batch['continuation_indices'][start:end][idx_min][
+                0]:batch['continuation_indices'][start:end][idx_min][-1] + 1]
+            metric_result_dict['context'].append(question)
+            metric_result_dict['correct_choice'].append(correct_choice)
+            metric_result_dict['selected_choice'].append(selected_choice)
+
             self.total += torch.tensor(1.0)
 
+        return metric_result_dict
+
     def compute(self):
         assert isinstance(self.correct, Tensor)
         assert isinstance(self.total, Tensor)
@@ -609,6 +683,8 @@ def __init__(self, dist_sync_on_step: bool = False):
         if self.eval_device is not None:
             self.eval_device = self.eval_device.upper()
 
+        self.metric_result_dict = {'context': [], 'output': [], 'result': [], 'sample_id': []}
+
     def get_client(self) -> EvalClient:
         """Returns a client for the appropriate remote platform."""
         client = None
@@ -647,7 +723,7 @@ def estimator(self, n: int, c: int, k: int) -> float:
             return 1.0
         return 1.0 - float(np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))
 
-    def _initialize_state(self, batch: dict[str, Any]):
+    def _initialize_state(self, batch: Dict[str, Any]):
         device = batch['input_ids'].device
         self.dataset_size = batch['dataset_size']
         self.pass_at_k = batch['pass_at_k']
@@ -689,15 +765,19 @@ def update(self, batch: Dict[str, Any], outputs: List[str], labels: List[str]):
         del labels  # never used
         client = self.get_client()
 
+        metric_result_dict = copy.deepcopy(self.metric_result_dict)
         for sample_id, code_gen, sample_prompt, test_inputs, test_outputs, entry_point, language in zip(
                 batch['sample_id'], outputs, batch['prompts'], batch['test_inputs'], batch['test_outputs'],
                 batch['entry_points'], batch['languages']):
 
             idx = sample_id
             self.total[idx] += 1.0
+            metric_result_dict['sample_id'].append(sample_id)
 
             code_gen = re.split(r'\n[A-Za-z0-9#`]', code_gen)[0]  # remove everything after function ends
             final_code = sample_prompt + code_gen  # combine prompt with the code generation
+            metric_result_dict['context'].append(sample_prompt)
+            metric_result_dict['output'].append(code_gen)
 
             test_results = []
             for test_input, test_output in zip(test_inputs, test_outputs):
@@ -714,8 +794,12 @@ def update(self, batch: Dict[str, Any], outputs: List[str], labels: List[str]):
 
             if all(test_results):
                 self.correct[idx] += 1.0
+                metric_result_dict['result'].append(1)
+            else:
+                metric_result_dict['result'].append(0)
 
         client.close()  # pyright: ignore [reportOptionalMemberAccess]
+        return metric_result_dict
 
     def compute(self):
         assert isinstance(self.correct, Tensor)

composer/models/huggingface.py

@@ -530,11 +530,17 @@ def get_metrics(self, is_train: bool = False) -> Dict[str, Metric]:
 
         return metrics if metrics else {}
 
-    def update_metric(self, batch: Any, outputs: Any, metric: Metric) -> None:
+    def update_metric(self, batch: Any, outputs: Any, metric: Metric) -> Dict:
         if getattr(metric, 'needs_batch', False):
-            metric.update(batch=batch, outputs=outputs, labels=self.labels)
+            metric_result = metric.update(batch=batch, outputs=outputs, labels=self.labels)
         else:
-            metric.update(outputs, self.labels)
+            metric_result = metric.update(outputs, self.labels)
+        if metric_result is not None:
+            # Add the metric name once for each datapoint in the batch
+            metric_result['metric_name'] = [metric.__class__.__name__ for _ in range(0, batch['input_ids'].shape[0])]
+        else:
+            metric_result = {}
+        return metric_result
 
     def get_metadata(self):
         model_output = {}

composer/trainer/trainer.py

@@ -3031,11 +3031,12 @@ def _eval_loop(
                                     outputs = self.state.outputs
 
                                 for metric in metrics.values():
-                                    self._original_model.update_metric(
+                                    metric_outputs = self._original_model.update_metric(
                                         self.state.batch,
                                         outputs,
                                         metric,
                                     )
+                                    self.state.metric_outputs = metric_outputs or {}
 
                     except RuntimeError as e:
                         if evaluator.auto_microbatching and _is_cuda_oom(e):