Fix repetition_penalty not activated issue (#9989)

This commit is contained in:
Yina Chen 2024-01-25 10:40:41 +08:00 committed by GitHub
parent 3bc3d0bbcd
commit 27338540c3

View file

@ -57,7 +57,7 @@ def generate(
new_speculative_kwargs = {}
for var in ['max_new_tokens', 'max_step_draft', 'th_stop_draft', 'do_sample',
'top_k', 'top_p', 'temperature', 'hf_adjust',
'auto_th_stop_draft', 'auto_parameters']:
'auto_th_stop_draft', 'auto_parameters', 'repetition_penalty']:
value = kwargs.pop(var, None)
if value is not None:
new_speculative_kwargs[var] = value