Use ungated models for unit tests (#196)

* Fix unit test

* Fix chat template tests

* Remove deprecated test

* up
This commit is contained in:
lewtun
2024-08-19 12:33:11 +02:00
committed by GitHub
parent 28bf902477
commit a8dcde2cd3
3 changed files with 11 additions and 25 deletions
-14
View File
@@ -15,7 +15,6 @@
import unittest
import torch
from transformers import AutoTokenizer
from alignment import (
DataArguments,
@@ -64,19 +63,6 @@ class GetTokenizerTest(unittest.TestCase):
tokenizer = get_tokenizer(self.model_args, DataArguments())
self.assertEqual(tokenizer.chat_template, DEFAULT_CHAT_TEMPLATE)
def test_default_chat_template_no_overwrite(self):
"""
If no chat template is passed explicitly in the config, then for models with a
`default_chat_template` but no `chat_template` we do not set a `chat_template`,
and that we do not change `default_chat_template`
"""
model_args = ModelArguments(model_name_or_path="m-a-p/OpenCodeInterpreter-SC2-7B")
base_tokenizer = AutoTokenizer.from_pretrained("m-a-p/OpenCodeInterpreter-SC2-7B")
processed_tokenizer = get_tokenizer(model_args, DataArguments())
assert getattr(processed_tokenizer, "chat_template") is None
self.assertEqual(base_tokenizer.default_chat_template, processed_tokenizer.default_chat_template)
def test_chatml_chat_template(self):
chat_template = "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
tokenizer = get_tokenizer(self.model_args, DataArguments(chat_template=chat_template))