diff --git a/litgpt/tokenizer.py b/litgpt/tokenizer.py index 2a3028de58..cace23dfbd 100644 --- a/litgpt/tokenizer.py +++ b/litgpt/tokenizer.py @@ -73,7 +73,7 @@ def __init__(self, checkpoint_dir: Path | str) -> None: self.apply_decoding_fix = None if (config_path := checkpoint_dir / "tokenizer_config.json").is_file(): with open(config_path, encoding="utf-8") as fp: - self.apply_decoding_fix = "LlamaTokenizer" in json.load(fp)["tokenizer_class"] + self.apply_decoding_fix = "LlamaTokenizer" in (json.load(fp).get("tokenizer_class") or "") @property def vocab_size(self) -> int: diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index c950de6043..95b52423f7 100644 --- a/tests/test_tokenizer.py +++ b/tests/test_tokenizer.py @@ -1,4 +1,5 @@ # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file. +import json import os import shutil from pathlib import Path @@ -112,6 +113,16 @@ def test_tokenizer_input_validation(): Tokenizer("cocofruit") +def test_tokenizer_config_without_tokenizer_class(tmp_path): + HFTokenizer(BPE()).save(str(tmp_path / "tokenizer.json")) + with open(tmp_path / "tokenizer_config.json", "w", encoding="utf-8") as fp: + json.dump({"add_bos_token": False}, fp) + + tokenizer = Tokenizer(tmp_path) + + assert tokenizer.apply_decoding_fix is False + + @pytest.mark.parametrize("use_bos_by_default", (True, False)) @pytest.mark.parametrize("encode_use_bos", (None, True, False)) @pytest.mark.parametrize("encode_use_eos", (True, False))