SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Prophetnet optimization (#9453)

Browse Source 
* Vectorized `ngram_attention_bias` calculation

* updated formatting with black

* Further optimization

* one (last) optimization
This commit is contained in:
guillaume-be
2021-01-07 11:41:58 +01:00
committed by GitHub
parent 28d74872cc
commit 390cf16bc8
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/transformers/models/prophetnet/modeling_prophetnet.py
View File

@@ -171,13 +171,15 @@ def ngram_attention_bias(sequence_length, ngram, device, dtype):
""" """
This function computes the bias for the predict stream This function computes the bias for the predict stream
""" """
bias = torch.ones((ngram, sequence_length, 2 * sequence_length), device=device, dtype=dtype) * float("-inf") left_block = torch.ones((ngram, sequence_length, sequence_length), device=device, dtype=dtype) * float("-inf")
right_block = left_block.detach().clone()
# create bias # create bias
for stream_idx in range(ngram): for stream_idx in range(ngram):
for i in range(sequence_length): right_block[stream_idx].fill_diagonal_(0, wrap=False)
bias[stream_idx, i, sequence_length + i] = 0 left_block[stream_idx].triu_(-stream_idx + 1)
bias[stream_idx, i, : max(i - stream_idx, 0) + 1] = 0
return bias left_block[:, :, 0] = 0
return torch.cat([left_block, right_block], dim=2)
def compute_relative_buckets(num_buckets, max_distance, relative_positions, is_bidirectional=False): def compute_relative_buckets(num_buckets, max_distance, relative_positions, is_bidirectional=False):