[WIP] Error logging callback

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
@@ -34,6 +35,7 @@
     'CheckpointSaver',
     'MLPerfCallback',
     'EarlyStopper',
+    'EvalOutputLogging',
     'ExportForInferenceCallback',
     'ThresholdStopper',
     'ImageVisualizer',

composer/callbacks/eval_output_logging_callback.py

@@ -0,0 +1,148 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Log model outputs and expected outputs during ICL evaluation."""
+
+import hashlib
+import os
+import random
+import shutil
+import time
+from typing import Callable, Optional
+
+from torch.utils.data import DataLoader
+
+from composer.core import Callback, State
+from composer.datasets.in_context_learning_evaluation import (InContextLearningCodeEvalDataset,
+                                                              InContextLearningLMTaskDataset,
+                                                              InContextLearningMultipleChoiceTaskDataset,
+                                                              InContextLearningQATaskDataset,
+                                                              InContextLearningSchemaTaskDataset)
+from composer.loggers import Logger
+from composer.loggers.console_logger import ConsoleLogger
+from composer.utils import MissingConditionalImportError, dist, maybe_create_object_store_from_uri, parse_uri
+
+ICLDatasetTypes = (InContextLearningLMTaskDataset, InContextLearningQATaskDataset,
+                   InContextLearningMultipleChoiceTaskDataset, InContextLearningSchemaTaskDataset,
+                   InContextLearningCodeEvalDataset)
+
+
+def _write(destination_path, src_file):
+    obj_store = maybe_create_object_store_from_uri(destination_path)
+    _, _, save_path = parse_uri(destination_path)
+    if obj_store is not None:
+        obj_store.upload_object(object_name=save_path, filename=src_file)
+    else:
+        with dist.local_rank_zero_download_and_wait(destination_path):
+            if dist.get_local_rank() == 0:
+                shutil.copy(src_file, destination_path)
+
+
+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 also responsible for providing a method for rendering the cached responses as strings.
+    This callback then accesses each eval benchmark during eval_end, retrieves the cached results,
+    and renders and and logs them in tabular format.
+
+    If subset_sample > 0, then only `subset_sample` of the outputs will be logged.
+
+    output_directory indicates where to write the tsv results, either can be a local directory or a cloud storage directory.
+    """
+
+    def __init__(self, subset_sample: int = -1, output_directory: Optional[str] = None):
+        self.subset_sample = subset_sample
+        self.table = {}
+        self.output_directory = output_directory if output_directory else os.getcwd()
+        self.hash = hashlib.sha256()
+        self.destination_file = None
+
+    def _write_tables_to_output_dir(self, state: State):
+        try:
+            import pandas as pd
+        except ImportError as e:
+            raise MissingConditionalImportError(extra_deps_group='pandas',
+                                                conda_package='pandas',
+                                                conda_channel='conda-forge') from e
+        # write tmp files
+        self.hash.update((str(time.time()) + str(random.randint(0, 1_000_000))).encode('utf-8'))
+        tmp_dir = os.getcwd() + '/' + self.hash.hexdigest()
+
+        if not os.path.exists(tmp_dir):
+            with dist.local_rank_zero_download_and_wait(tmp_dir):
+                if dist.get_local_rank() == 0:
+                    os.mkdir(tmp_dir)
+
+        full_df = pd.DataFrame()
+        file_name = f'eval-outputs-ba{state.timestamp.batch.value}.tsv'
+
+        for benchmark in self.table:
+            cols, rows = self.table[benchmark]
+            rows = [[e.encode('unicode_escape') if isinstance(e, str) else e for e in row] for row in rows]
+            df = pd.DataFrame.from_records(data=rows, columns=cols)
+            df['benchmark'] = benchmark
+            full_df = pd.concat([full_df, df], ignore_index=True)
+
+        with dist.local_rank_zero_download_and_wait(f'{tmp_dir}/{file_name}'):
+            if dist.get_local_rank() == 0:
+                with open(f'{tmp_dir}/{file_name}', 'wb') as f:
+                    full_df.to_csv(f, sep='\t', index=False)
+
+        # copy/upload tmp files
+        _write(destination_path=f'{self.output_directory}/{file_name}', src_file=f'{tmp_dir}/{file_name}')
+        os.remove(f'{tmp_dir}/{file_name}')
+        self.destination_file = f'{self.output_directory}/{file_name}'
+
+        # delete tmp files
+        os.rmdir(tmp_dir)
+
+    def _prep_response_cache(self, state, cache):
+        benchmark = state.dataloader_label
+        for metric in state.eval_metrics[benchmark].values():
+            if hasattr(metric, 'reset_response_cache'):
+                metric.reset_response_cache(cache)
+
+    def eval_start(self, state: State, logger: Logger) -> None:
+        # eval start runs before each benchmark's evaluator (either in training or eval)
+        self._prep_response_cache(state, True)
+
+    def eval_after_all(self, state: State, logger: Logger) -> None:
+        # eval after all runs after all evaluators have completed during eval within training
+        #  (either in training or eval)
+        self._write_tables_to_output_dir(state)
+        self.table = {}
+
+    def eval_standalone_end(self, state: State, logger: Logger) -> None:
+        # eval standalone end runs after all evaluators have completed during a direct call to trainer.eval()
+        self._write_tables_to_output_dir(state)
+        self.table = {}
+
+    def eval_end(self, state: State, logger: Logger) -> None:
+        # eval start runs after each benchmark's evaluator
+        # during each eval, only a single dataloader/benchmark will be active
+        assert state.dataloader is not None
+        assert isinstance(state.dataloader, DataLoader)
+        if hasattr(state.dataloader, 'dataset') and isinstance(state.dataloader.dataset, ICLDatasetTypes):
+            assert isinstance(state.dataloader.dataset, ICLDatasetTypes)
+            if hasattr(state.dataloader.dataset, 'tokenizer'):
+                tokenizer = state.dataloader.dataset.tokenizer
+                benchmark = state.dataloader_label
+                assert benchmark is not None
+                assert isinstance(benchmark, str)
+                for metric_name, metric in state.eval_metrics[benchmark].items():
+                    if hasattr(metric, 'format_response_cache'):
+                        assert isinstance(metric.format_response_cache, Callable)
+                        format_response_cache: Callable = metric.format_response_cache
+                        columns, rows = format_response_cache(tokenizer)
+
+                        if columns is not None and rows is not None:
+                            if self.subset_sample > 0:
+                                rows = random.sample(rows, min(len(rows), self.subset_sample))
+                            for destination in logger.destinations:
+                                if not isinstance(destination, ConsoleLogger):
+                                    # don't log to console because it will pollute the console too much
+                                    destination.log_table(columns, rows, f'icl_outputs/{benchmark}/{metric_name}')
+
+                            self.table[f'{benchmark}_{metric_name}'] = (columns, rows)
+        self._prep_response_cache(state, False)

composer/datasets/in_context_learning_evaluation.py

@@ -164,6 +164,7 @@ def __init__(
             if dist.get_local_rank() == 0:
                 get_file(dataset_uri, destination_path, overwrite=True)
         dataset = load_dataset('json', data_files=destination_path, split='train', streaming=False)
+
         self.samples = self._read_dataset(dataset)
         self.samples = strip_data(self.samples)
         self.tokenizer = tokenizer
@@ -225,6 +226,7 @@ def _prep_examples(self,
             example_delimiter (str): The delimiter used to separate each individual context/continuation pair
             continuation_delimiter (str): The delimiter used to separate each context from its continuation
             question_prelimiter (str): The text to prepend to each question
+            cot_delimiter (str): The delimiter used to separate the chain-of-thought (if present) from the final model response.
             fewshot_rng (random.Random): Random number generator to use for fewshot sampling
             cot_delimiter (str): The delimiter used to separate the chain-of-thought (if present) from the final model response.
 
@@ -241,7 +243,6 @@ def _prep_examples(self,
             prompt_and_fewshot = self._format_prompt_and_fewshot(num_fewshot, prompt_string, example_delimiter,
                                                                  continuation_delimiter, question_prelimiter,
                                                                  cot_delimiter, fewshot_rng, sample_idx)
-
             ctxt = self.samples[sample_idx]['context']
             ctxt = f'{question_prelimiter}{ctxt}'
             if len(prompt_and_fewshot) > 0:
@@ -260,6 +261,7 @@ def _prep_examples(self,
                 encoded_example['preamble']['input_ids'] = encoded_example['preamble']['input_ids'][:-1]
 
             encoded_example['context'] = self.tokenizer(ctxt, add_special_tokens=False)
+
             encoded_example['aliases'] = list(self.samples[sample_idx]['aliases'])
             encoded_example['cot_delimiter'] = cot_delimiter
             examples.append(encoded_example)

composer/loggers/in_memory_logger.py

@@ -79,7 +79,9 @@ def log_table(self, columns: List[str], rows: List[List[Any]], name: str = 'Tabl
             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)
+        table = pd.DataFrame.from_records(data=rows, columns=columns).to_json(orient='split',
+                                                                              index=False,
+                                                                              force_ascii=False)
         self.tables[name] = table
 
     def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None: