Skip to content

Commit

Permalink
Update GenerationDefaults and comments
Browse files Browse the repository at this point in the history
  • Loading branch information
mgoin authored Oct 23, 2023
1 parent 97e6c6c commit f1e644d
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions src/deepsparse/transformers/pipelines/text_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,16 +62,21 @@
__all__ = ["TextGenerationPipeline"]


# Based off of https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig # noqa E501
class GenerationDefaults:
num_return_sequences = 1
# Parameters that control the length of the output
max_length = None
max_new_tokens = 100
output_scores = False
top_k = 0
top_p = 0.0
repetition_penalty = 0.0
# Parameters that control the generation strategy used
do_sample = False
# Parameters for manipulation of the model output logits
temperature = 1.0
top_k = 50
top_p = 1.0
repetition_penalty = 1.0
# Parameters that define the outputs
num_return_sequences = 1
output_scores = False


class FinishReason(Enum):
Expand Down

0 comments on commit f1e644d

Please sign in to comment.