From 4531aa2c2c7fc2c4fb8de007a48631cfec38a07b Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Wed, 31 Jan 2024 14:56:03 +0100 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com> --- optimum/neuron/generation/token_selector.py | 2 +- optimum/neuron/modeling.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/neuron/generation/token_selector.py b/optimum/neuron/generation/token_selector.py index 62a95fb1c..2d64d03fc 100644 --- a/optimum/neuron/generation/token_selector.py +++ b/optimum/neuron/generation/token_selector.py @@ -71,7 +71,7 @@ def create( The model provides the internal helpers allowing to select the logits processors and stopping criterias. max_seq_length (`int`): The maximum number of input + generated tokens for this model. It depends on the model compilation parameters. - stopping_criteria (`Optional[transformers.generation.StoppingCriteriaList]): + stopping_criteria (`Optional[transformers.generation.StoppingCriteriaList], defaults to `None`): Custom stopping criteria that complement the default stopping criteria built from arguments and a generation config. Return: diff --git a/optimum/neuron/modeling.py b/optimum/neuron/modeling.py index 5b6944ea5..9f37f9b9e 100644 --- a/optimum/neuron/modeling.py +++ b/optimum/neuron/modeling.py @@ -749,7 +749,7 @@ def generate( priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model configuration. Please note that unspecified parameters will inherit [`~transformers.generation.GenerationConfig`]'s default values, whose documentation should be checked to parameterize generation. - stopping_criteria (`Optional[transformers.generation.StoppingCriteriaList]): + stopping_criteria (`Optional[transformers.generation.StoppingCriteriaList], defaults to `None`): Custom stopping criteria that complement the default stopping criteria built from arguments and a generation config.