remove token_type_ids from TokenizerBertGeneration output (#10070)

huggingface · Feb 8, 2021 · 84acf0c · 84acf0c
1 parent e4bf991
commit 84acf0c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/bert_generation/tokenization_bert_generation.py b/src/transformers/models/bert_generation/tokenization_bert_generation.py
@@ -58,6 +58,7 @@ class BertGenerationTokenizer(PreTrainedTokenizer):
     pretrained_vocab_files_map = {"vocab_file": {"bert_for_seq_generation": tokenizer_url}}
     max_model_input_sizes = {"bert_for_seq_generation": 512}
     prefix_tokens: List[int] = []
+    model_input_names = ["input_ids", "attention_mask"]
 
     def __init__(
         self,