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

make RepetitionPenaltyLogitsProcessor faster (#9600)

Browse Source
This commit is contained in:
LSinev
2021-01-20 12:23:01 +03:00
committed by GitHub
parent a1ad16a446
commit a98173cc45
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/transformers/generation_logits_process.py
View File

@@ -155,13 +155,12 @@ class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
self.penalty = penalty self.penalty = penalty
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
ranges = torch.arange(scores.shape[0]) score = torch.gather(scores, 1, input_ids)
score = scores[ranges[:, None], input_ids]
# if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability # if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability
score = torch.where(score < 0, score * self.penalty, score / self.penalty) score = torch.where(score < 0, score * self.penalty, score / self.penalty)
scores[ranges[:, None], input_ids] = score scores.scatter_(1, input_ids, score)
return scores return scores