Add custom stopping criteria to ICL generate tasks

composer/datasets/in_context_learning_evaluation.py

@@ -16,6 +16,7 @@
 
 from composer.core import DataSpec
 from composer.core.data_spec import _default_split_batch, _split_list
+from composer.datasets.utils import stop_sequences_criteria
 from composer.utils import MissingConditionalImportError, dist, get_file
 
 if TYPE_CHECKING:
@@ -139,21 +140,21 @@ def _read_dataset(self, dataset: Dataset) -> List[Dict[str, str]]:
             })
         return result
 
-    def __init__(
-        self,
-        dataset_uri: str,
-        tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast],
-        max_seq_len: int,
-        pad_tok_id: int,
-        num_fewshot: int,
-        prompt_string: str,
-        example_delimiter: str,
-        continuation_delimiter: str,
-        destination_path: str,
-        question_prelimiter: str,
-        fewshot_random_seed: int,
-        cot_delimiter: str = '',
-    ):
+    def __init__(self,
+                 dataset_uri: str,
+                 tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast],
+                 max_seq_len: int,
+                 pad_tok_id: int,
+                 num_fewshot: int,
+                 prompt_string: str,
+                 example_delimiter: str,
+                 continuation_delimiter: str,
+                 destination_path: str,
+                 question_prelimiter: str,
+                 fewshot_random_seed: int,
+                 cot_delimiter: str = '',
+                 early_stopping_criteria: Optional[List[str]] = None,
+                 do_normalization: bool = True):
         if tokenizer.eos_token_id is None:
             raise ValueError('`InContextLearningQATaskDataset` tokenizer must have non-null `eos_token_id`')
         try:
@@ -166,6 +167,8 @@ 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.early_stopping_criteria = early_stopping_criteria
+        self.do_normalization = do_normalization
         self.samples = self._read_dataset(dataset)
         self.samples = strip_data(self.samples)
         self.tokenizer = tokenizer
@@ -299,17 +302,22 @@ def collate_fn(self, data):
             # We will search for the answer within the portion of the model response
             # beginning with `cot_delimiter`
             cot_delimiter = sample['cot_delimiter']
-
+        stopping_criteria = None
+        if self.early_stopping_criteria:
+            stopping_criteria = stop_sequences_criteria(self.tokenizer, self.early_stopping_criteria, len(inputs))
         batch = {
             'input_ids': torch.stack(inputs),
             'mode': 'generate',
             'labels': answers,
             'cot_delimiter': cot_delimiter,
             'generation_length': self.max_answer_length,
+            'stopping_criteria': self.early_stopping_criteria,
+            'do_normalization': self.do_normalization,
             'generation_kwargs': {
                 'pad_token_id': self.pad_tok_id,
                 'use_cache': True,
-                'eos_token_id': self.tokenizer.eos_token_id
+                'stopping_criteria': stopping_criteria,
+                'eos_token_id': self.tokenizer.eos_token_id,
             }
         }
 
@@ -323,7 +331,9 @@ def split_batch(self, batch: Any, microbatch_size: int):
         # Don't split kwargs that don't change
         # Normally split torch tensors
         # List split lists of strings
-        no_split = ['mode', 'generation_length', 'generation_kwargs', 'cot_delimiter']
+        no_split = [
+            'mode', 'generation_length', 'generation_kwargs', 'cot_delimiter', 'do_normalization', 'stopping_criteria'
+        ]
         normal_split = ['input_ids', 'attention_mask']
         list_split = ['labels']
         chunked = {}
@@ -339,7 +349,7 @@ def split_batch(self, batch: Any, microbatch_size: int):
                 raise ValueError(f'Unexpected key {k}')
         num_chunks = len(chunked['input_ids'])
         for k, v in batch.items():
-            if isinstance(v, (int, float, str, bool, dict)):
+            if k in no_split:
                 chunked[k] = [v] * num_chunks
         return [{k: v[idx] for k, v in chunked.items()} for idx in range(num_chunks)]
 
@@ -950,6 +960,7 @@ def __init__(
         pass_at_k: int = 1,
         top_p: Optional[float] = 0.95,
         top_k: Optional[int] = 40,
+        early_stopping_criteria: Optional[List[str]] = None,
     ):
         if tokenizer.eos_token_id is None:
             raise ValueError('`InContextLearningCodeEvalDataset` tokenizer must have non-null `eos_token_id`')
@@ -963,6 +974,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.early_stopping_criteria = early_stopping_criteria
         self.samples = list(
             dataset.map(
                 lambda examples: {
@@ -1100,6 +1112,9 @@ def collate_fn(self, data):
             test_outputs.append(test_output)
             languages.append(language)
 
+        stopping_criteria = None
+        if self.early_stopping_criteria:
+            stopping_criteria = stop_sequences_criteria(self.tokenizer, self.early_stopping_criteria, len(inputs))
         batch = {
             'input_ids': torch.stack(inputs),
             'mode': 'generate',
@@ -1121,6 +1136,7 @@ def collate_fn(self, data):
                 'top_p': self.top_p,
                 'top_k': self.top_k,
                 'use_cache': True,
+                'stopping_criteria': stopping_criteria,
                 'eos_token_id': self.tokenizer.eos_token_id
             }
         }
@@ -1161,23 +1177,24 @@ def split_batch(self, batch: Any, microbatch_size: int):
 
 
 def build_icl_dataloader(
-    icl_task_type: str,
-    dataset_uri: str,
-    tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast],
-    batch_size: int,
-    max_seq_len: int,
-    pad_tok_id: int,
-    num_fewshot: int,
-    prompt_string: str,  # e.g. 'translate english to french:'
-    example_delimiter: str,  # e.g. '\n'
-    continuation_delimiter: str,  # e.g. ''
-    destination_path: str,
-    question_prelimiter: str = '',  # e.g. 'Question: '
-    cot_delimiter: str = '',
-    fewshot_random_seed: int = 1234,
-    pass_at_k: int = 1,
-    generations_per_sample: int = 1,
-) -> DataSpec:
+        icl_task_type: str,
+        dataset_uri: str,
+        tokenizer: Union[transformers.PreTrainedTokenizer, transformers.PreTrainedTokenizerFast],
+        batch_size: int,
+        max_seq_len: int,
+        pad_tok_id: int,
+        num_fewshot: int,
+        prompt_string: str,  # e.g. 'translate english to french:'
+        example_delimiter: str,  # e.g. '\n'
+        continuation_delimiter: str,  # e.g. ''
+        destination_path: str,
+        question_prelimiter: str = '',  # e.g. 'Question: '
+        cot_delimiter: str = '',
+        fewshot_random_seed: int = 1234,
+        pass_at_k: int = 1,
+        generations_per_sample: int = 1,
+        early_stopping_criteria: Optional[List[str]] = None,
+        do_normalization: bool = True) -> DataSpec:
     if icl_task_type == 'multiple_choice':
         dataset = InContextLearningMultipleChoiceTaskDataset(dataset_uri,
                                                              tokenizer,
@@ -1228,7 +1245,9 @@ def build_icl_dataloader(
                                                  destination_path=destination_path,
                                                  question_prelimiter=question_prelimiter,
                                                  fewshot_random_seed=fewshot_random_seed,
-                                                 cot_delimiter=cot_delimiter)
+                                                 cot_delimiter=cot_delimiter,
+                                                 early_stopping_criteria=early_stopping_criteria,
+                                                 do_normalization=do_normalization)
         effective_batchsize = batch_size
     elif icl_task_type == 'code_evaluation':
         dataset = InContextLearningCodeEvalDataset(dataset_uri,
@@ -1242,7 +1261,8 @@ def build_icl_dataloader(
                                                    code_prelimiter=question_prelimiter,
                                                    fewshot_random_seed=fewshot_random_seed,
                                                    pass_at_k=pass_at_k,
-                                                   generations_per_sample=generations_per_sample)
+                                                   generations_per_sample=generations_per_sample,
+                                                   early_stopping_criteria=early_stopping_criteria)
         effective_batchsize = batch_size
     else:
         raise Exception(f'Unrecognized ICL task type: {icl_task_type}')
@@ -1328,7 +1348,9 @@ def get_icl_task_dataloader(
         pass_at_k: int = 1,
         generations_per_sample: int = 1,
         cot_delimiter: str = '',
-        has_categories: bool = False) -> Union[DataSpec, Dict[str, DataSpec]]:
+        has_categories: bool = False,
+        early_stopping_criteria: Optional[List[str]] = None,
+        do_normalization: bool = True) -> Union[DataSpec, Dict[str, DataSpec]]:
     """This constructs a dataloader (or dataloaders if has_categories is True) capable of evaluating LLMs on in-context learning language modeling tasks, for example LAMBADA. An example usage is below:
 
     >>> dl = get_icl_task_dataloader(
@@ -1381,41 +1403,41 @@ def get_icl_task_dataloader(
         categories = sorted(output_files.keys())
         for category in categories:
             partition_uri = output_files[category]
-            result_dls[category] = build_icl_dataloader(
-                icl_task_type,
-                partition_uri,
-                tokenizer,
-                batch_size,
-                max_seq_len,
-                pad_tok_id,
-                num_fewshot,
-                prompt_string,
-                example_delimiter,
-                continuation_delimiter,
-                partition_uri + '_tmp',
-                question_prelimiter,
-                cot_delimiter,
-                fewshot_random_seed,
-                pass_at_k,
-                generations_per_sample,
-            )
+            result_dls[category] = build_icl_dataloader(icl_task_type,
+                                                        partition_uri,
+                                                        tokenizer,
+                                                        batch_size,
+                                                        max_seq_len,
+                                                        pad_tok_id,
+                                                        num_fewshot,
+                                                        prompt_string,
+                                                        example_delimiter,
+                                                        continuation_delimiter,
+                                                        partition_uri + '_tmp',
+                                                        question_prelimiter,
+                                                        cot_delimiter,
+                                                        fewshot_random_seed,
+                                                        pass_at_k,
+                                                        generations_per_sample,
+                                                        early_stopping_criteria=early_stopping_criteria,
+                                                        do_normalization=do_normalization)
         return result_dls
     else:
-        return build_icl_dataloader(
-            icl_task_type,
-            dataset_uri,
-            tokenizer,
-            batch_size,
-            max_seq_len,
-            pad_tok_id,
-            num_fewshot,
-            prompt_string,
-            example_delimiter,
-            continuation_delimiter,
-            destination_path,
-            question_prelimiter,
-            cot_delimiter,
-            fewshot_random_seed,
-            pass_at_k,
-            generations_per_sample,
-        )
+        return build_icl_dataloader(icl_task_type,
+                                    dataset_uri,
+                                    tokenizer,
+                                    batch_size,
+                                    max_seq_len,
+                                    pad_tok_id,
+                                    num_fewshot,
+                                    prompt_string,
+                                    example_delimiter,
+                                    continuation_delimiter,
+                                    destination_path,
+                                    question_prelimiter,
+                                    cot_delimiter,
+                                    fewshot_random_seed,
+                                    pass_at_k,
+                                    generations_per_sample,
+                                    early_stopping_criteria=early_stopping_criteria,
+                                    do_normalization=do_normalization)

composer/datasets/utils.py

@@ -6,7 +6,7 @@
 import logging
 import textwrap
 import warnings
-from typing import Callable, List, Tuple, Union
+from typing import TYPE_CHECKING, Callable, List, Tuple, Union
 
 import numpy as np
 import torch
@@ -15,6 +15,10 @@
 from torchvision.datasets import VisionDataset
 
 from composer.core import Batch
+from composer.utils import MissingConditionalImportError
+
+if TYPE_CHECKING:
+    import transformers
 
 __all__ = [
     'add_vision_dataset_transform',
@@ -166,3 +170,63 @@ def add_vision_dataset_transform(dataset: VisionDataset, transform: Callable, is
         else:
             dataset.transform = transforms.Compose([dataset.transform, transform])
             log.warning(transform_added_logstring)
+
+
+try:
+    import transformers
+    del transformers  # unused
+except ImportError as e:
+    raise MissingConditionalImportError(extra_deps_group='nlp',
+                                        conda_package='transformers',
+                                        conda_channel='conda-forge') from e
+
+
+class MultiTokenEOSCriteria(transformers.StoppingCriteria):
+    """Criteria to stop on the specified multi-token sequence.
+    Slightly modified from: https://github.com/EleutherAI/lm-evaluation-harness/blob/78545d42f2ca95c6fe0ed220d456eeb94f4485e9/lm_eval/utils.py#L614-L649
+    """
+
+    def __init__(
+        self,
+        sequence: str,
+        tokenizer: transformers.PreTrainedTokenizer,
+        batch_size: int,
+    ) -> None:
+        self.done_tracker = [False] * batch_size
+        self.sequence = sequence
+        self.sequence_ids = tokenizer.encode(sequence, add_special_tokens=False)
+        # we look back for 2 more tokens than it takes to encode our stop sequence
+        # because tokenizers suck, and a model might generate `['\n', '\n']` but our `sequence` is `['\n\n']`
+        # and we don't want to mistakenly not stop a generation because our
+        # (string) stop sequence was output in a different tokenization
+
+        # NOTE: there is a minor danger that this will end up looking back 2 tokens into the past, into the inputs to the model,
+        # and stopping generation immediately as a result. With only 2 extra tokens of lookback, this risk is minimized
+        self.sequence_id_len = len(self.sequence_ids) + 2
+        self.tokenizer = tokenizer
+
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        # For efficiency, we compare the last n tokens where n is the number of tokens in the stop_sequence
+        lookback_ids_batch = input_ids[:, :][:, -self.sequence_id_len:]
+
+        lookback_tokens_batch = self.tokenizer.batch_decode(lookback_ids_batch)
+        for i, done in enumerate(self.done_tracker):
+            if i >= len(lookback_tokens_batch):
+                # the last batch of a dataset may b smaller than `batch_size`
+                # automatically set those indices in the done_tracker to True
+                # since those indices don't show up in the current batch
+                self.done_tracker[i] = True
+                break
+            elif not done:
+                self.done_tracker[i] = self.sequence in lookback_tokens_batch[i]
+        return False not in self.done_tracker
+
+
+def stop_sequences_criteria(
+    tokenizer: transformers.PreTrainedTokenizer,
+    stop_sequences: List[str],
+    batch_size: int,
+) -> transformers.StoppingCriteriaList:
+    return transformers.StoppingCriteriaList([
+        *[MultiTokenEOSCriteria(sequence, tokenizer, batch_size) for sequence in stop_sequences],
+    ])

composer/metrics/nlp.py

@@ -270,13 +270,23 @@ def update(self, outputs: List[str], labels: List[List[str]], batch: Optional[Di
         if batch is None:
             batch = {}
         cot_delimiter = batch.get('cot_delimiter', '')
+        do_normalization = batch.get('do_normalization', True)
+        stopping_criteria = batch.get('stopping_criteria', None)
         for sample_output, sample_labels in zip(outputs, labels):
             final_answer = sample_output
+
+            if stopping_criteria is not None and len(stopping_criteria) > 0:
+                final_answer = re.split('|'.join(stopping_criteria), final_answer)[0]
+
             if cot_delimiter is not None and len(cot_delimiter) > 0:
                 final_answer = final_answer.split(cot_delimiter)[-1]
 
-            cleaned_final_answer = self.normalize_answer(final_answer)
-            cleaned_sample_labels = {self.normalize_answer(label) for label in sample_labels}
+            if do_normalization:
+                cleaned_final_answer = self.normalize_answer(final_answer)
+                cleaned_sample_labels = {self.normalize_answer(label) for label in sample_labels}
+            else:
+                cleaned_final_answer = final_answer
+                cleaned_sample_labels = set(sample_labels)
 
             if any(cleaned_final_answer.startswith(label) for label in cleaned_sample_labels):
                 self.correct += torch.tensor(1.0)