Skip to content

Commit

Permalink
Add back in generation types (huggingface#28681)
Browse files Browse the repository at this point in the history
  • Loading branch information
@amyeroberts @DavidAfonsoValente
amyeroberts authored and DavidAfonsoValente committed Jan 25, 2024
1 parent 2c0fa70 commit 8b3f4b8
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/transformers/generation/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,11 @@ class GenerateBeamEncoderDecoderOutput(ModelOutput):
BeamSearchEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput
BeamSampleEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput

GreedySearchOutput = Union[GreedySearchEncoderDecoderOutput, GreedySearchDecoderOnlyOutput]
SampleOutput = Union[SampleEncoderDecoderOutput, SampleDecoderOnlyOutput]
BeamSearchOutput = Union[BeamSearchEncoderDecoderOutput, BeamSearchDecoderOnlyOutput]
BeamSampleOutput = Union[BeamSampleEncoderDecoderOutput, BeamSampleDecoderOnlyOutput]
ContrastiveSearchOutput = Union[ContrastiveSearchEncoderDecoderOutput, ContrastiveSearchDecoderOnlyOutput]

# Typing shortcuts
GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput]
Expand Down

0 comments on commit 8b3f4b8

Please sign in to comment.