Prevent unnecessarily overwriting the default Hugging Face chat template (#2183)

Signed-off-by: Keshav Santhanam <[email protected]>

Author: santhnm2

megatron/core/tokenizers/text/libraries/huggingface_tokenizer.py

@@ -69,15 +69,13 @@ def __init__(
                     pretrained_model_name_or_path=tokenizer_path,
                     use_fast=use_fast,
                     trust_remote_code=trust_remote_code,
-                    chat_template=chat_template,
                 )
             elif merges_file is None:
                 self.tokenizer = AutoTokenizer.from_pretrained(
                     pretrained_model_name_or_path=tokenizer_path,
                     vocab_file=vocab_file,
                     use_fast=use_fast,
                     trust_remote_code=trust_remote_code,
-                    chat_template=chat_template,
                 )
             else:
                 self.tokenizer = AutoTokenizer.from_pretrained(
@@ -86,14 +84,21 @@ def __init__(
                     merge_files=merges_file,
                     use_fast=use_fast,
                     trust_remote_code=trust_remote_code,
-                    chat_template=chat_template,
                 )
         except Exception as e:
             raise ValueError(
                 'Unable to instantiate HuggingFace AutoTokenizer '
                 f'for {tokenizer_path}. Exception: {e}'
             )
 
+        # Store the tokenizer's existing chat template if the user does not provide
+        # a custom chat template. Otherwise, override the default chat template with
+        # the user-provided template.
+        if chat_template is None:
+            chat_template = self.tokenizer.chat_template
+        else:
+            self.tokenizer.chat_template = chat_template
+
         self.include_special_tokens = include_special_tokens
         self.original_vocab_size = len(self.tokenizer)
         self.chat_template = chat_template

megatron/core/tokenizers/text/text_tokenizer.py

@@ -37,13 +37,17 @@ def __init__(self, path: str, config: dict, **kwargs) -> None:
         self._tokenizer = self._restore_model(**kwargs)
         self.additional_args = kwargs
         self.path = path
-        if (
-            config.get("chat_template", None) is None
-            and kwargs.get("chat_template", None) is not None
-        ):
-            self.chat_template = kwargs.get("chat_template", None)
+
+        config_template = config.get("chat_template", None)
+        tokenizer_template = getattr(self._tokenizer, "chat_template", None)
+        kwargs_template = kwargs.get("chat_template", None)
+
+        if config_template is not None:
+            self.chat_template = config_template
+        elif tokenizer_template is not None:
+            self.chat_template = tokenizer_template
         else:
-            self.chat_template = config.get("chat_template", None)
+            self.chat_template = kwargs_template
 
     def _restore_model(self, **kwargs) -> MegatronTokenizerTextAbstract:
         """Returns tokenizer library object."""