setup training

9f75565e · Rémi Louf · Julien Chaumond · 4735c2af · 9f75565e · 9f75565e
Commit 9f75565e authored Nov 08, 2019 by Rémi Louf Committed by Julien Chaumond Dec 09, 2019
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 3 deletions

requirements.txt requirements.txt +0 -2

transformers/generate/beam_search.py transformers/generate/beam_search.py +1 -1

No files found.
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,5 +10,3 @@ regex
 sentencepiece
 # For XLM
 sacremoses
-# For ROUGE
-pyrouge
--- a/transformers/generate/beam_search.py
+++ b/transformers/generate/beam_search.py
@@ -166,7 +166,7 @@ class BeamSearch(object):
        for step in range(self.max_length):

            decoder_input = fit_to_block_size(self.growing_beams, block_size, self.pad_token_id)
-            kwargs_decoder["attention_mask"] = build_mask(decoder_input)
+            kwargs_decoder["attention_mask"] = build_mask(decoder_input, self.pad_token_id)
            outputs = self.model.decoder(decoder_input, **kwargs_decoder)

            next_token_scores = outputs[0][:, -1, :].squeeze(1)