Stream ModelOutputs

[dmarx]

What does this PR do?

The "streaming" feature for text generation is currently constrained to streaming token ids. If a user requests an output dict, the stream still only generates token ids without other attributes that may have been requested such as scores, raw logits, activations. After the stream has been consumed, the function returns its final output, which will be the originally requested output dict.

This PR aligns the return type of the streamer with the requested return type by encapsulating the logic that determines how the return value is constructed. In so doing, this change also permits users to stream richer representations than just token ids.

Summary of changes:

• Adds a ._prepare_output() function to GenerationMixin, encapsulates output formatting logic
• ._prepare_output() replaces nested conditionals that previously built return values (beam samplers excluded)
• Streamer.put() receives its input from ._prepare_output() invocation rather than restricted to token IDs tensor
• Inputs to Streamer.put() are no longer sent to the CPU prior to being passed to .put()
  • When desired, delegates responsibility for moving the tensor to the provided Streamer.
• Adds OutputStreamer and OutputIteratorStreamer classes

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Did you write any new necessary tests?
  • OutputIteratorStreamerTester to be modularized with fixtures, DRYed
• test against hydra-node
• fix inconsistent list output from OutputStreamer.on_ready
• flesh out rest of intermediate return values
  • scores
  • logits
  • attention
  • hidden
• add some kind of "output formatter" factory to ensure outputs are uniform across the class w/o repeating a ton of logic.
  • encapsulate output construction logic in a function
  • refactor to use function
  • add tests validating output type consistency and parity with streamer.put()
• ensure tests fully cover targeted cases
  • I'm concerned that injecting bad data directly into these attributes doesn't seem to trigger test failures
  • encoder attentions
  • cross attentions
  • [contrastive search] decoder attention
  • suspect tests not covering assisted decoding
  • added encoder and encoder-decoder models to test suite, need to smoke test
• update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
  • Send the CW documentation team a CR
  • docstrings
  • typehints
  • explicit input args for _prepare_output() (? maybe it's fine as is?)
• fix import
• fix inconsistent tensor dimensions
• fix sigterm on hydra-node
• [ ] revisit ECHO behavior?
  • on our end, just skipping over streamed outputs when output.score is None
  • should be sufficient for HF's end if we just document that the ECHO response will have null scores/logits/activations
• reimplement their TextIteratorStreamer using the new OutputIteraterSreamer

Who can review?

(Hidden @ tag, uncomment when ready to send to HF)

[dmarx]

moved to upstream: huggingface#29545