Fix Bloom test

huggingface · Jul 2, 2024 · e84aa4f · e84aa4f
1 parent 5e146b2
commit e84aa4f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/tests/models/bloom/test_tokenization_bloom.py b/tests/models/bloom/test_tokenization_bloom.py
@@ -135,6 +135,7 @@ def test_encodings_from_xnli_dataset(self):
     @require_jinja
     def test_tokenization_for_chat(self):
         tokenizer = self.get_rust_tokenizer()
+        tokenizer.chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
         test_chats = [
             [{"role": "system", "content": "You are a helpful chatbot."}, {"role": "user", "content": "Hello!"}],
             [

diff --git a/tests/models/gpt2/test_tokenization_gpt2.py b/tests/models/gpt2/test_tokenization_gpt2.py
@@ -280,7 +280,7 @@ def test_special_tokens_mask_input_pairs_and_bos_token(self):
     @require_jinja
     def test_tokenization_for_chat(self):
         tokenizer = GPT2Tokenizer.from_pretrained(self.tmpdirname)
-        tokenizer.chat_template = '{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}'
+        tokenizer.chat_template = "{% for message in messages %}{{ message.content }}{{ eos_token }}{% endfor %}"
         test_chats = [
             [{"role": "system", "content": "You are a helpful chatbot."}, {"role": "user", "content": "Hello!"}],
             [