Unverified Commit aeef4823 authored by Thomas Wolf's avatar Thomas Wolf Committed by GitHub
Browse files

Merge pull request #2303 from patrickvonplaten/fix_error_with_repetition_penalty

fix repetition penalty error in modeling_utils.py
parents 0412f3d9 18e5bdbe
......@@ -728,6 +728,10 @@ class PreTrainedModel(nn.Module):
if repetition_penalty != 1.0:
for i in range(batch_size):
for previous_tokens in set(input_ids[i].tolist()):
# if score < 0 then repetition penalty has to multiplied to reduce the previous token probability
if next_token_logits[i, previous_tokens] < 0:
next_token_logits[i, previous_tokens] *= repetition_penalty
else:
next_token_logits[i, previous_tokens] /= repetition_penalty
if do_sample:
......@@ -807,6 +811,10 @@ class PreTrainedModel(nn.Module):
if repetition_penalty != 1.0:
for i in range(batch_size * num_beams):
for previous_tokens in set(input_ids[i].tolist()):
# if score < 0 then repetition penalty has to multiplied to reduce the previous token probability
if scores[i, previous_tokens] < 0:
scores[i, previous_tokens] *= repetition_penalty
else:
scores[i, previous_tokens] /= repetition_penalty
if do_sample:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment