Flax Whisper predicts erroneous exclamation mark
Author: sanchit-gandhi
Pull request: no

### System Info

- `transformers` version: 4.27.0.dev0
- Platform: Linux-4.19.0-21-cloud-amd64-x86_64-with-debian-10.13
- Python version: 3.7.12
- Huggingface_hub version: 0.12.1
- PyTorch version (GPU?): 1.12.1+cpu (False)
- Tensorflow version (GPU?): 2.11.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.6.4 (gpu)
- Jax version: 0.3.25
- JaxLib version: 0.3.25
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@sanchit-gandhi @andyehrenberg @ArthurZucker 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Code snippet:
```python
from transformers import FlaxWhisperForConditionalGeneration, WhisperProcessor
from datasets import load_dataset

model = FlaxWhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny")
processor = WhisperProcessor.from_pretrained("openai/whisper-tiny")

librispeech = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
sample = librispeech[0]["audio"]["array"]

input_features = processor(sample, return_tensors="np").input_features

pred_ids = model.generate(input_features)
pred_text = processor.batch_decode(pred_ids.sequences)
print(pred_text)
```
**Print Output:**
```
['
!

 Mischekvilder is the apostle of the middle classes, and we are glad to welcome his gospel.
']
```

We see an extra `!` after the `
` token that should't be there.

Do you fancy taking a look into this one @andyehrenberg? Otherwise can try and find time next week.

### Expected behavior

Should return:
```
['


 Mischekvilder is the apostle of the middle classes, and we are glad to welcome his gospel.
']
```

Comments:
I think the expected behavior should be returning:

`['



 Mischekvilder is the apostle of the middle classes, and we are glad to welcome his gospel.
']`

The generation config has [[1, None], ...] for forced decoder ids, so the language is predicted by the model, and for some reason it's sampling "!". Maybe we should have a processor for only sampling from the language tokens for that first step when it is supposed to predict the language.

Indeed - that's correct we should be seeing the language token predicted at the second index!

Just checked and `!` is the zero-th token in the tokenizer vocab -> maybe somethings going astray with the forced tokens logits processor when a `None` is passed as the forced token?

I'm seeing the problem - `force_token_array.at[index].set(token)` when `token` is `None` sets the value at `index` to 0. We should just make it so when `token` is None, we keep the value at that index at -1.

So should update FlaxForceTokensLogitsProcessor to:

```
def __init__(self, force_token_map):
    force_token_map = dict(force_token_map)
    # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
    # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
    # Indexes without forced tokens will have a negative value.
    force_token_array = jnp.ones((max(force_token_map.keys()) + 1), dtype=jnp.int32) * -1
    for index, token in force_token_map.items():
        if token is not None:
            force_token_array = force_token_array.at[index].set(token)
    self.force_token_array = jnp.int32(force_token_array)
```

Also, just keep in mind that `forced_decoder_ids` has to be a static argument for jitted functions. The workaround I use is having empty `forced_decoder_ids` and instead passing them into `decoder_input_ids` when I know the forced ids might change.

Nice one @andyehrenberg! That must indeed be the root cause of the problem ✅. We'll have to pass the forced decoder ids as static argnums when we `pmap` the generate function in #21764

Hi
Could you please specify how to set forced_decoder_ids for the FlaxWhisperForConditionalGeneration object?
@sanchit-gandhi 

You should either modify the ` model.generation_config.forced_decoder_ids`  or when calling `generate`, set the `language`, `task`  and `return_timestamps` arguments. You can also pass them as `decoder_input_ids`  (which is also an argument of the `generate()`  function or ` forced_decoder_ids`. 

> You should either modify the ` model.generation_config.forced_decoder_ids` or when calling `generate`, set the `language`, `task` and `return_timestamps` arguments. You can also pass them as `decoder_input_ids` (which is also an argument of the `generate()` function or ` forced_decoder_ids`.

It does not respect the forced_decoder_ids when I pass it to model.generation_config.
This is my code:
```
model = FlaxWhisperForConditionalGeneration.from_pretrained(model_id, dtype=jnp.float16, from_pt=True)
jit_generate = jax.jit(model.generate, static_argnames=["max_length"])
model.generation_config.forced_decoder_ids = processor.get_decoder_prompt_ids(language="en", task="translate")
pred_ids = jit_generate(input_features, max_length=128)
```

But the result is not translation to english (as same as when forced_decoder_ids is set to None)


Also, when I set `language` to `generate()` function, it raises error:
```
pred_ids = jit_generate(input_features, max_length=128, language="
")

File /opt/conda/lib/python3.8/site-packages/jax/_src/api_util.py:568, in _str_abstractify(x)
    567 def _str_abstractify(x):
--> 568   raise TypeError(f"Argument '{x}' of type {type(x)} is not a valid JAX type")

TypeError: Argument '
' of type <class 'str'> is not a valid JAX type
```

I've also tested `language='en'` and `language='english'` and the result is the same (following error)

That is because you are not using the latest version of transformers. All of this was adressed in #21965

Your error regarding strings not being a valid JAX type can be fixed by setting the language prior to compiling and keeping it static. It also looks like you’re changing the model’s generation config after wrapping its generate method in jit, which could be causing problems. My guidance is the set your generation parameters how you want, and then get a `partial(model.generate, arg1=val2, …)` and then compile that function (or just use static argnames).

It is resolved by passing the `language` parameter in the static_argnames:
```
jit_generate = jax.jit(model.generate, static_argnames=["max_length", "language"])
input_features = jnp.array(input_features, dtype=jnp.float16)
pred_ids = jit_generate(input_features, max_length=128, language='
')
```
Thanks @andyehrenberg  and @ArthurZucker