past_key_values supporting more-than-one-token inputs
Author: namespace-Pt
Pull request: no

### System Info

- `transformers` version: 4.30.0
- Platform: Linux-5.4.0-113-generic-x86_64-with-glibc2.31
- Python version: 3.10.11
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

`past_key_values` does not work when my `input_ids` are more than 1 token. For example:
```python
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

inputs = "I love hugging face. Me too"
inputs = tokenizer(inputs, return_tensors="pt")
inputs1 = {}
inputs2 = {}
for k, v in inputs.items():
    inputs1[k] = v[:, :-3]
    inputs2[k] = v[:, -3:]

outputs = model(**inputs)
input1_outputs = model(**inputs1)
# Error!!
input2_outputs = model(**inputs2, past_key_values=input1_outputs.past_key_values)
``` 

### Expected behavior

I find the error is because you only extend the keys and values [here](https://github.com/huggingface/transformers/blob/fe861e578f50dc9c06de33cd361d2f625017e624/src/transformers/models/gpt2/modeling_gpt2.py#L319), while you forget to extend the `attention_mask` to the same size as the extended keys and values and hence the error.

I also find this is common across different models, e.g. gpt-neo, gpt2. I think this is an urgent problem because many downstream applications like chatbots require this feature.

I think you can extend the attention mask simply by concatenating `torch.ones((batch_size, past_length))` in front of the input attention mask to solve the problem. Here is my work around:

```python
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

inputs = "I love hugging face. Me too"
inputs = tokenizer(inputs, return_tensors="pt")
inputs1 = {}
inputs2 = {}
for k, v in inputs.items():
    inputs1[k] = v[:, :-3]
    inputs2[k] = v[:, -3:]

outputs = model(**inputs)
input1_outputs = model(**inputs1)

# without the following line, will raise errors
inputs2["attention_mask"] = inputs.attention_mask

input2_outputs = model(**inputs2, past_key_values=input1_outputs.past_key_values)

# check
print(((input2_outputs.logits - outputs.logits[:, -3:]) < 1e-4).all())
```



Comments:
Hey! Thanks for opening an issue! 🤗 
The reason why this is not working currently is that GPT2 is a pretty old model, thus it requires the attention mask to be passed when you want to generate. 
Now this is not an issue most of the time people use `gpt2_model.generate(input_ids, attention_mask)` and thus don't need to handle the past key values on their own! This is why, no it's not a very urgent problem and it's pretty much expected. Someone had a similar question see [here.](https://github.com/huggingface/transformers/issues/16811)

The issue rather lies in the creation of the positional ids, see in #18104 

Hey! Thanks for replying so soon.

I tried `.generate` method with `input_ids` and `past_key_values` but it does not work as expected when I have more than one token in the `input_ids`. 

To be specific, assume I'm building a QA system with GPT2. My input would be like 
```
Q: a question\nA: an answer\nQ: a new question\nA:
```
After generating the first answer, I have `past_key_values` until the token `answer`. However, when I want to generate the second answer, due to the insertion of the new question and prompts, I have to input `\nQ: a new question\nA: ` together with the `past_key_values` to the model. 

I expect the model to compute the hidden states of the given input, then generate the next token. However, according to the [code](https://github.com/huggingface/transformers/blob/35eac0df75c692c5b93c12f7eaf3279cab8bd7ce/src/transformers/models/gpt2/modeling_gpt2.py#L1012), the model will automatically truncate the `input_ids` as long as there is `past_key_values` passed alongside. This leads to false generation results.

Here is my snippet:
```
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

inputs = "I love you. Me too"
inputs = tokenizer(inputs, return_tensors="pt")
inputs1 = {}
inputs2 = {}
inputs3 = {}
inputs4 = {}
for k, v in inputs.items():
    inputs1[k] = v[:, :-3]
    inputs2[k] = v[:, -3:]
    inputs3[k] = v[:, :-1]
    inputs4[k] = v[:, -1:]

print(f"All inputs: {tokenizer.batch_decode(inputs['input_ids'])}")
print(f"Inputs1: {tokenizer.batch_decode(inputs1['input_ids'])}")
print(f"Inputs2: {tokenizer.batch_decode(inputs2['input_ids'])}")
print(f"Inputs3: {tokenizer.batch_decode(inputs3['input_ids'])}")
print(f"Inputs4: {tokenizer.batch_decode(inputs4['input_ids'])}")

# 1. Generate without cache. This is the standard output.
outputs = model.generate(**inputs, max_new_tokens=5)
print(tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:]))

# 2. WRONG!!! Generate with partial past_key_values. Extend attention_mask by past_length because the model expect the input_ids of shape [B, 1] when past_key_values is not None
outputs1 = model(**inputs1)
past_length = outputs1.past_key_values[0][0].size(-2)
inputs2["attention_mask"] = torch.cat([torch.ones(1, past_length), inputs2['attention_mask'][:, :1]], dim=-1)
print(inputs2, past_length)
outputs2 = model.generate(**inputs2, past_key_values=outputs1.past_key_values, max_new_tokens=5)
print(tokenizer.batch_decode(outputs2[:, inputs2['input_ids'].shape[1]:]))

# 3. CORRECT!!! Generate with past_key_values of all previous tokens except the most recent one. Extend attention_mask by past_length because the model expect the input_ids of shape [B, 1] when past_key_values is not None
outputs3 = model(**inputs3)
past_length = outputs3.past_key_values[0][0].size(-2)
inputs4["attention_mask"] = torch.cat([torch.ones(1, past_length), inputs4['attention_mask'][:, :1]], dim=-1)
print(inputs4, past_length)
outputs4 = model.generate(**inputs4, past_key_values=outputs3.past_key_values, max_new_tokens=5)
print(tokenizer.batch_decode(outputs4[:, inputs4['input_ids'].shape[1]:]))
```

What you are trying to do is very akin to the [`QuestionAnsweringPipeline](https://huggingface.co/docs/transformers/main/main_classes/pipelines#transformers.QuestionAnsweringPipeline)`, which implements all the pre-processing and post processing.

The generate function is made for general generation, and uses indeed the `prepare_inputs_for_generation` function. What you are expecting is not a supported behaviour, but rather a specific usage. In the `generate` function, we expect the `new_tokens` to be a single token per batch:

https://github.com/huggingface/transformers/blob/f092997ca669750d4f32ada127b2624bd450aee5/src/transformers/generation/utils.py#L2474

 It should be possible to hack your way trough the code, probably by writing a logits processor that only returns the last prediction, so that when you compute the next token: 

https://github.com/huggingface/transformers/blob/f092997ca669750d4f32ada127b2624bd450aee5/src/transformers/generation/utils.py#L2465

the shape is correct. You also need to modify the `prepare_inputs_for_generation`.

cc @gante for visibility


Got it. Thank you Arthur.