Keep special token ids the tokenizer does not define

src/transformers/trainer_utils.py

@@ -1169,6 +1169,35 @@ def compare_trainer_and_checkpoint_args(training_args, trainer_state):
         logger.warning_once(warning_str)
 
 
+def _align_special_token(name, tokenizer, config, generation_config, updated_tokens, several=False):
+    """
+    Aligns one special token id held by the configs with the tokenizer, recording the change in `updated_tokens`. A
+    token the tokenizer does not define is left untouched. With `several`, the generation config keeps the tokens it
+    already holds alongside the tokenizer's one, as a list.
+    """
+    tokenizer_token_id = getattr(tokenizer, name)
+    if tokenizer_token_id is None:
+        return
+    generation_token_ids = getattr(generation_config, name) if generation_config is not None else None
+    if not isinstance(generation_token_ids, list):
+        generation_token_ids = [generation_token_ids]
+
+    if tokenizer_token_id == getattr(config, name, None) and (
+        generation_config is None or tokenizer_token_id in generation_token_ids
+    ):
+        return
+
+    updated_tokens[name] = tokenizer_token_id
+    setattr(config, name, tokenizer_token_id)
+    if generation_config is None:
+        return
+    if several:
+        # Any of the tokens listed here will halt generation, so the ones already defined are preserved.
+        setattr(generation_config, name, [tokenizer_token_id, *(t for t in generation_token_ids if t is not None)])
+    else:
+        setattr(generation_config, name, tokenizer_token_id)
+
+
 def align_special_tokens(model, processing_class):
     """
     Aligns the special tokens of the tokenizer with the model configs.
@@ -1189,54 +1218,17 @@ def align_special_tokens(model, processing_class):
     model_has_generation_config = hasattr(model, "generation_config") and model.generation_config is not None
     updated_tokens = {}
 
-    # 1 - Align EOS token. EOS is more complex than the others, as `generation_config` may hold more than one EOS
-    # token.
-    tokenizer_has_new_eos = tokenizer.eos_token_id != getattr(config, "eos_token_id", None)
-    if model_has_generation_config:
-        # `generation_config.eos_token_id` is None: direct comparison
-        if model.generation_config.eos_token_id is None:
-            tokenizer_has_new_eos |= tokenizer.eos_token_id != model.generation_config.eos_token_id
-        else:
-            # `generation_config.eos_token_id` is an `int`: convert it to list (and continue below)
-            if isinstance(model.generation_config.eos_token_id, int):
-                model.generation_config.eos_token_id = [model.generation_config.eos_token_id]
-            # `generation_config.eos_token_id` is a `list`: check if the tokenizer's EOS token is in the list
-            tokenizer_has_new_eos |= tokenizer.eos_token_id not in model.generation_config.eos_token_id
-
-    if tokenizer_has_new_eos:
-        updated_tokens["eos_token_id"] = tokenizer.eos_token_id
-        config.eos_token_id = tokenizer.eos_token_id
-        # The generation config may hold more than one EOS token. We preserve the original EOS tokens: any of the
-        # EOS tokens defined here will halt generation.
-        if model_has_generation_config:
-            all_eos_tokens = [tokenizer.eos_token_id]
-            if model.generation_config.eos_token_id is not None:
-                all_eos_tokens += list(model.generation_config.eos_token_id)
-            model.generation_config.eos_token_id = [token for token in all_eos_tokens if token is not None]
-
-    # 2 - Align BOS
-    tokenizer_has_new_bos = tokenizer.bos_token_id != getattr(config, "bos_token_id", None)
-    if model_has_generation_config:
-        tokenizer_has_new_bos |= tokenizer.bos_token_id != model.generation_config.bos_token_id
-
-    if tokenizer_has_new_bos:
-        updated_tokens["bos_token_id"] = tokenizer.bos_token_id
-        config.bos_token_id = tokenizer.bos_token_id
-        if model_has_generation_config:
-            model.generation_config.bos_token_id = tokenizer.bos_token_id
-
-    # 3 - Align PAD
-    tokenizer_has_new_pad = tokenizer.pad_token_id != getattr(config, "pad_token_id", None)
-    if model_has_generation_config:
-        tokenizer_has_new_pad |= tokenizer.pad_token_id != model.generation_config.pad_token_id
-
-    if tokenizer_has_new_pad:
-        updated_tokens["pad_token_id"] = tokenizer.pad_token_id
-        config.pad_token_id = tokenizer.pad_token_id
-        if model_has_generation_config:
-            model.generation_config.pad_token_id = tokenizer.pad_token_id
-
-    # 4 - Warn users about the changes
+    # A token the tokenizer does not define is not a statement that the model should drop the value the config
+    # declares: `_special_tokens_map` is initialized with `None` for every special token, so an absent attribute
+    # and a deliberately cleared one are indistinguishable. Only values the tokenizer defines are propagated.
+
+    # 1 - Align the tokens. The generation config may hold more than one EOS token.
+    generation_config = model.generation_config if model_has_generation_config else None
+    _align_special_token("eos_token_id", tokenizer, config, generation_config, updated_tokens, several=True)
+    _align_special_token("bos_token_id", tokenizer, config, generation_config, updated_tokens)
+    _align_special_token("pad_token_id", tokenizer, config, generation_config, updated_tokens)
+
+    # 2 - Warn users about the changes
     if len(updated_tokens) > 0:
         logger.warning(
             "The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. "

tests/trainer/test_trainer.py

@@ -1365,6 +1365,25 @@ def test_special_token_alignment_composite_config(self):
         self.assertEqual(text_config.bos_token_id, tokenizer.bos_token_id)
         self.assertEqual(text_config.pad_token_id, tokenizer.pad_token_id)
 
+    def test_special_token_alignment_keeps_tokens_the_tokenizer_does_not_define(self):
+        """
+        Tests that a special token the tokenizer does not define is left untouched on the model configs, rather than
+        being removed from them. `_special_tokens_map` is initialized with `None` for every special token, so a
+        tokenizer that never declared one is indistinguishable from a tokenizer whose token was deliberately cleared.
+        Dropping the value in that case would silently remove an id the checkpoint declares.
+        """
+        model = AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM")
+        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-LlamaForCausalLM")
+
+        model.config.bos_token_id = 1
+        model.generation_config.bos_token_id = 1
+        tokenizer.bos_token = None
+
+        align_special_tokens(model, tokenizer)
+
+        self.assertEqual(model.config.bos_token_id, 1)
+        self.assertEqual(model.generation_config.bos_token_id, 1)
+
     def test_trainer_works_without_model_config(self):
         """
         Tests that models without a `config` parameter can still be trained.