[TF Longformer] Improve Speed for TF Longformer (#6447)

* add tf graph compile tests * fix conflict * remove more tf transpose statements * fix conflicts * fix comment typos * move function to class function * fix black * fix black * make style
huggingface · Aug 26, 2020 · 858b7d5 · 858b7d5
1 parent a75c64d
commit 858b7d5
Show file tree

Hide file tree

Showing 5 changed files with 257 additions and 134 deletions.
diff --git a/src/transformers/modeling_tf_bert.py b/src/transformers/modeling_tf_bert.py
@@ -1088,7 +1088,7 @@ def __init__(self, config, *inputs, **kwargs):
             config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
         )
 
-    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(
         tokenizer_class=_TOKENIZER_FOR_DOC,
         checkpoint="bert-base-cased",

diff --git a/src/transformers/modeling_tf_electra.py b/src/transformers/modeling_tf_electra.py
@@ -677,7 +677,7 @@ def __init__(self, config, *inputs, **kwargs):
         self.electra = TFElectraMainLayer(config, name="electra")
         self.classifier = TFElectraClassificationHead(config, name="classifier")
 
-    @add_start_docstrings_to_callable(ELECTRA_INPUTS_DOCSTRING.format("(batch_size, num_choices, sequence_length)"))
+    @add_start_docstrings_to_callable(ELECTRA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
     @add_code_sample_docstrings(
         tokenizer_class=_TOKENIZER_FOR_DOC,
         checkpoint="google/electra-small-discriminator",