diff --git a/auto_round/calib_dataset.py b/auto_round/calib_dataset.py index 7d2e3a2f8..c52e4c901 100644 --- a/auto_round/calib_dataset.py +++ b/auto_round/calib_dataset.py @@ -359,7 +359,7 @@ def get_ultrachat_dataset( def is_instruct_tokenizer(tokenizer): try: - out = tokenizer.apply_chat_template([{"role": "user", "content": "Hi"}]) + out = tokenizer.apply_chat_template([{"role": "user", "content": "Hi"}], tokenize=False) return bool(out and len(out) > 0) except Exception: return False @@ -371,7 +371,7 @@ def is_instruct_tokenizer(tokenizer): apply_chat_template = True elif not is_instruct and apply_chat_template: logger.info("Tokenizer is not an instruct/chat model, but apply_chat_template=True. Setting to False.") - apply_chat_template = False + apply_chat_template = False def tokenize_example_batch(examples): if not apply_chat_template: diff --git a/test/test_cpu/utils/test_calib_dataset.py b/test/test_cpu/utils/test_calib_dataset.py index adc129f8d..478c30f38 100644 --- a/test/test_cpu/utils/test_calib_dataset.py +++ b/test/test_cpu/utils/test_calib_dataset.py @@ -4,9 +4,11 @@ import pytest import torch +from transformers import BatchEncoding from transformers import AutoModelForCausalLM, AutoTokenizer from auto_round import AutoRound +from auto_round.calib_dataset import get_ultrachat_dataset from ...helpers import get_model_path, opt_name_or_path @@ -110,3 +112,64 @@ def test_combine_dataset2(self, tiny_opt_model_path): nsamples=1, ) autoround.quantize() + + +class _FakeStreamingDataset: + def __init__(self): + self.mapped_output = None + + def shuffle(self, seed=42): + return self + + def take(self, count): + return self + + def map(self, fn, batched=True): + self.mapped_output = fn({"messages": [[{"role": "user", "content": "hello"}]]}) + return self + + +class _FakeV5ChatTokenizer: + def __init__(self): + self.rendered_messages = None + + def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False, **kwargs): + if tokenize: + return BatchEncoding({"input_ids": [[1, 2]], "attention_mask": [[1, 1]]}) + return f"templated::{messages[-1]['content']}" + + def __call__(self, texts, truncation=True, max_length=None): + self.rendered_messages = texts + return { + "input_ids": [[len(text)] for text in texts], + "attention_mask": [[1] for _ in texts], + } + + +class _FakePlainTokenizer(_FakeV5ChatTokenizer): + def apply_chat_template(self, messages, tokenize=True, add_generation_prompt=False, **kwargs): + raise ValueError("chat templates are not supported") + + +def test_ultrachat_dataset_keeps_chat_template_for_v5_tokenizers(monkeypatch): + tokenizer = _FakeV5ChatTokenizer() + dataset = _FakeStreamingDataset() + monkeypatch.setattr("auto_round.calib_dataset.load_dataset", lambda *args, **kwargs: dataset) + + result = get_ultrachat_dataset(tokenizer=tokenizer, seqlen=128, apply_chat_template=True) + + assert result is dataset + assert tokenizer.rendered_messages == ["templated::hello"] + assert dataset.mapped_output["input_ids"] == [[16]] + + +def test_ultrachat_dataset_disables_chat_template_for_plain_tokenizers(monkeypatch): + tokenizer = _FakePlainTokenizer() + dataset = _FakeStreamingDataset() + monkeypatch.setattr("auto_round.calib_dataset.load_dataset", lambda *args, **kwargs: dataset) + + result = get_ultrachat_dataset(tokenizer=tokenizer, seqlen=128, apply_chat_template=True) + + assert result is dataset + assert tokenizer.rendered_messages == ["hello"] + assert dataset.mapped_output["input_ids"] == [[5]]