Ensure correct padding token for Phi and Pythia models (#3899)

ludwig-ai · Jan 18, 2024 · ea50811 · ea50811
1 parent 6e93ea7
commit ea50811
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/ludwig/utils/tokenizers.py b/ludwig/utils/tokenizers.py
@@ -845,11 +845,16 @@ def get_unk_token(self) -> str:
     def _set_pad_token(self) -> None:
         """Sets the pad token and pad token ID for the tokenizer."""
 
+        # CodeGenTokenizer Used by Phi-2
+        # GPTNeoXTokenizerFast Used by Pythia
         from transformers import (
+            CodeGenTokenizer,
+            CodeGenTokenizerFast,
             CodeLlamaTokenizer,
             CodeLlamaTokenizerFast,
             GPT2Tokenizer,
             GPT2TokenizerFast,
+            GPTNeoXTokenizerFast,
             LlamaTokenizer,
             LlamaTokenizerFast,
         )
@@ -865,12 +870,15 @@ def _set_pad_token(self) -> None:
         if any(
             isinstance(self.tokenizer, t)
             for t in [
+                CodeGenTokenizer,
+                CodeGenTokenizerFast,
+                CodeLlamaTokenizer,
+                CodeLlamaTokenizerFast,
                 GPT2Tokenizer,
                 GPT2TokenizerFast,
+                GPTNeoXTokenizerFast,
                 LlamaTokenizer,
                 LlamaTokenizerFast,
-                CodeLlamaTokenizer,
-                CodeLlamaTokenizerFast,
             ]
         ):
             if hasattr(self.tokenizer, "eos_token") and self.tokenizer.eos_token is not None: