Bart can make decoder_input_ids from labels (huggingface#6758)

fabiocapsouza · Nov 15, 2020 · f53d5ad · f53d5ad
1 parent 07d0a53
commit f53d5ad
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/modeling_bart.py b/src/transformers/modeling_bart.py
@@ -58,8 +58,8 @@
     "facebook/bart-large-cnn",
     "facebook/bart-large-xsum",
     "facebook/mbart-large-en-ro",
-    # See all BART models at https://huggingface.co/models?filter=bart
 ]
+# This list is incomplete. See all BART models at https://huggingface.co/models?filter=bart
 
 
 BART_START_DOCSTRING = r"""
@@ -1045,6 +1045,8 @@ def forward(
 
         if labels is not None:
             use_cache = False
+            if decoder_input_ids is None:
+                decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id)
 
         outputs = self.model(
             input_ids,