Add eot token to ICL generate kwargs

composer/datasets/in_context_learning_evaluation.py

@@ -306,7 +306,8 @@ def collate_fn(self, data):
             'generation_length': self.max_answer_length,
             'generation_kwargs': {
                 'pad_token_id': self.pad_tok_id,
-                'use_cache': True
+                'use_cache': True,
+                'eos_token_id': self.tokenizer.eos_token_id
             }
         }
 
@@ -1116,6 +1117,7 @@ def collate_fn(self, data):
                 'top_p': self.top_p,
                 'top_k': self.top_k,
                 'use_cache': True,
+                'eos_token_id': self.tokenizer.eos_token_id
             }
         }
         batch['attention_mask'] = ~(batch['input_ids'] == self.pad_tok_id)

tests/datasets/test_in_context_learning_datasets.py

@@ -545,6 +545,7 @@ def test_qa_task_dataloader(dataset_uri, tiny_gpt2_tokenizer, tmp_path, num_fews
     assert tuple(batch['attention_mask'].shape) == (batch_size, seqlen - maximum_answer_length)
     assert batch['mode'] == 'generate'
     # the maximum generation length from the small test data
+
     assert batch['generation_length'] == maximum_answer_length
     assert all(item[0] == tokenizer.eos_token_id for item in batch['input_ids'])
 
@@ -559,6 +560,7 @@ def test_qa_task_dataloader(dataset_uri, tiny_gpt2_tokenizer, tmp_path, num_fews
         for found, expected in zip(batch['labels'], [['David Seville'], ['Skorpio', 'Scorpio']]))
     assert decoded_batch[0].endswith('Q: Who was the man behind The Chipmunks?\nA:')
     assert decoded_batch[1].endswith('Q: What star sign is Jamie Lee Curtis?\nA:')
+    assert 'eos_token_id' in batch['generation_kwargs']
 
 
 @pytest.mark.parametrize('dataset_uri', ['gsm8k_small.jsonl'])