Optionally return past key values from generate
Author: dblakely
Pull request: no

### Feature request

The idea would be to optionally return `past_key_values` inside the generation objects (`SampleEncoderDecoderOutput`, etc). This could be controlled by a flag called `output_past_key_values` that's passed to `generate` and then forwarded to `sample`, etc.

### Motivation

Perhaps this is niche, but my team and I need often need to obtain the past keys and values when generating in order to manipulate them a bit and then feed them back in for subsequent calls to `generate`. We currently do this with a custom version of `sample`, but this results in us having to copy and paste a lot of code. Would it be possible to allow `past_key_values` to be optionally returned by `generate`?

### Your contribution

If you all approve of the feature idea, I'd be able to implement it and submit a PR. 

Comments:
Hello, I have the exactly same issue! Could you please share an implementation of yours that return past key value when using generate?

> Hello, I have the exactly same issue! Could you please share an implementation of yours that return past key value when using generate?

Hi, what we do is something like this:

```python
class CustomSampleMixin:
    """Have an HF model return past_key_values from generate by inheriting this mixin. 
    For example:
        ```
        class YourModel(CustomSampleMixin, T5ForConditionalGeneration)
        ```
    """

    def sample(self, all_the_normal_args, output_past_key_values):
        """Custom sample method that's mostly copied from Huggingface's generation_utils.sample method
        """
        ...

        past_key_values = None
        while True:
            # forward pass to get next token
            outputs = self(**model_inputs, return_dict=True)
            past_key_values = outputs.past_key_values
            ...

        if return_dict_in_generate:
            return CustomGenerationOutput(
                sequences=input_ids,
                scores=scores,
                decoder_attentions=decoder_attentions,
                cross_attentions=cross_attentions,
                decoder_hidden_states=decoder_hidden_states,
                past_key_values=past_key_values if output_past_key_values else None,
                encoder_outputs=model_kwargs["encoder_outputs"],
            )

        return input_ids
```

Then when we're using it:

```python
past_key_values = None
while True:
    model_kwargs["past"] = past_key_values # Note that HF calls it "past" in the `model_kwargs`
    outputs = model.generate(output_past_key_values=True, **model_kwargs)
    past_key_values = outputs.past_key_values
    # post-process outputs, post-process past_key_values
    ...
```

This approach works for our purposes but does mean that we need to copy and maintain a lot of extra code from Huggingface. Plus, if you want to get `past_key_values` from `beam_sample`, `greedy_search`, etc instead of just `sample`, you have to make custom versions of each of those as well.

Thank you for sharing!

Hi @patrickvonplaten, what do you think about this idea?

I'd actually be fine with adding this to main generate, maybe already by default as soon as `return_dict_in_generate` is set to True, not sure if we necessarily need a new `output_...` input arg. @gante what do you think?

@patil-suraj what do you think here?

I'm fine with returning `past_key_values` from `generate` since we already allow to return other model outputs like attentions and hidden_states.  And a new argument is not necessary IMO, since the model always returns past when `use_cache=True` (the default case), so a new argument to control this won't be necessary.

I agree, we can return `past_key_values` when `use_cache=True`. 

It will, however, be an API change (adds a field to the output, which is an `OrderedDict` subclass), so any user iterating over the output will be impacted. I suspect it is a very uncommon use case, and thus the utility of exposing `past_key_values` exceeds potential pain points. WDYT @patil-suraj @patrickvonplaten? 

If you agree, I can add this to my to-do list.

Think it's fine to extend the len of the tuple / `ModelOutput`, we don't consider this a breaking change. @patil-suraj do you want to give it a try to implement this? @gante you could then fully focus on finishing TF generate :heart_eyes: 

Will open a PR for it this week :) 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Associated PR hasn't really been merged. So this issue cannot be closed.

Hey folks 👋 

#25086 was merged.

If you install from `main` and add `return_dict_in_generate=True` to `generate`, `past_key_values` will be part of the output, assuming your model is configured with `use_cache=True` (the default).

You can then pass `past_key_values` to `generate` to continue generating!