Add early stopping for Bark generation via logits processor
Author: isaac-chung
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue) https://github.com/huggingface/transformers/issues/26672


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
@ylacombe maybe we can continue the [conversation](https://github.com/huggingface/transformers/issues/26672#issuecomment-1753060623) here.

> Ideally, we'd have another test on BarkSemanticModelTest, but I'm not sure how to proceed yet.
Do you have any ideas?

I'm not entirely sure. Maybe we could assert outputs from `self.model.generate` with the new arg somehow?

> could be possibly passed to BarkModel.generate kwargs without causing issues

To confirm that we support this, maybe we should add to `BarkModelIntegrationTests.test_generate_end_to_end_with_sub_models_args` as well?

Let's try to do both!

I think I managed to add to `BarkModelIntegrationTests` without issues. But I'd like to align on how to proceed with `BarkSemanticModelTest`. Specifically:
1. Only a few tests assert the outputs. As I'm unsure what to expect, I might print the outputs and assert those
2. I've been manually trying to fill in `BarkSemanticGenerationConfig` so that the `generate()` call does not fail. Not sure if there's a more efficient way.

@ylacombe thanks! Regarding 1, let's take `BarkModelIntegrationTests.test_generate_end_to_end_with_sub_models_args` for example, the test does not assert any outputs and it simply runs `.generate()`. Would that be fine here?

Let's try to find a case where the semantic model has to stop. You can get inspiration from that test: https://github.com/huggingface/transformers/blob/d085662c596066bad82942b1e15819f675cbd15e/tests/models/bark/test_modeling_bark.py#L904-L921

So basically, an example where, the same seed, the last output tokens are different, do you think it's possible?

If we set `min_eos_p` to anything that's non-zero, we only get the eos_token (set to 10000 for open-end generation). Here is what passed.
```
    @slow
    def test_generate_semantic_early_stop(self):
        input_ids = self.inputs

        # fmt: off
        # check first ids
        expected_output_ids = [10000]
        # fmt: on

        self.semantic_generation_config.min_eos_p = 0.05

        # greedy decoding
        with torch.no_grad():
            output_ids = self.model.semantic.generate(
                **input_ids,
                do_sample=False,
                temperature=1.0,
                semantic_generation_config=self.semantic_generation_config,
            )

        self.assertListEqual(output_ids[0, : len(expected_output_ids)].tolist(), expected_output_ids)
``` 
Is that what you have in mind?

Oh, that seems weird, have you tried with another generation strategy ? (i.e `do_sample=True, temperature=...`)? If you have the same results, it's probably on the logit processor side!

Yep getting the same output even after varying those parameters. In terms of expected outputs, 
- for this test, is it accurate that we want a list of similar output_ids, just potentially shorter?
- for the logit processor unit test, is the current expected output wrong?

Just wondering what would be a good way to start debugging this as audio generation is super new to me 🙏 

No worries here, let's do this step by step:
1. Is the logic correct?
2. If going from `min_eos_p` to `np.log(self.min_eos_p)` is correct? might be the wrong logic here, meaning that `probability(eos_token)` might be different from `torch.log(scores[:, eos_token_id]`. In fact I strongly suspect it is where we are wrong. is it log softmax or softmax? In anyways, we should make sure it is correct for every generation strategy
3. make sure `test_early_stop_processor` is correct 

If we still have this behavior after verifying those steps, we can think of other ways of veryfing it 

Good plan, thank you! 

Re: logic, seems like we can obtain log_prob from scores by
```python
logprobs = torch.nn.functional.log_softmax(scores.float(), dim=-1)
```
Then we can compare `log(min_eos_p)` directly.

Re: test_early_stop_processor, I realized that I had not implemented the suggestion:
> So basically, an example where, the same seed, the last output tokens are different, do you think it's possible?

I've done so now that I've read things through again :D . I do wonder if it will be flaky at times due to its stochastic nature?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26675). All of your documentation changes will be reflected on that endpoint.

Gentle nudge @ylacombe to let this bubble back up

Regarding using a stopping criteria, I don't think it's possible at the moment -> quoting #26672
> * Ideally, we'd add a [custom stopping criteria](https://huggingface.co/docs/transformers/v4.34.0/en/internal/generation_utils#transformers.StoppingCriteria), but as indicated in [custom stopping_critriea function doesn't receive logits scores (receives None instead) #23674](https://github.com/huggingface/transformers/issues/23674) it's not yet possible to use [`scores`](https://huggingface.co/docs/transformers/v4.34.0/en/internal/generation_utils#transformers.StoppingCriteria.__call__.scores) without setting `return_dict_in_generate=True, output_scores=True `.



It receives None because `output_scores` and `return_dict` are not properly set

Yes of course, but don't you think users should have the liberty to set  `output_scores` and `return_dict` as they want ?

For sure. So the goal here is by default to always stop early? (actually not returning the scores might be better in terms of memory ?) 
What I mean is that the stopping criterias are meant to be used that way 😉 

> For sure. So the goal here is by default to always stop early? (actually not returning the scores might be better in terms of memory ?) What I mean is that the stopping criterias are meant to be used that way 😉

Yes this is the goal here. Totally agree on the stopping criteria usage! Actually I haven't find a stopping criteria which uses `scores` yet, maybe because of the limitation of having to use `return_dict_in_generate=True, output_scores=True`.  #23674 is a discussion on this and I believe this is under @gante's radar! What do you recommend in the meantime ?

Hey @ylacombe / @ArthurZucker , please let me know if there's anything else I can do to further this PR.

btw, regarding it being a logits processor vs stopping criteria: it is my impression that we want to generate an EOS token under the conditions defined here. Since we want to generate a token, it has to be a logits processor.

(the main difference between them is that the stopping criteria stops generation right away, and doesn't add any new token -- for batched generation, this can make a big difference)

@ylacombe I've run this command and all tests are passing ✅ 
```
RUN_SLOW=yes python -m unittest tests.models.bark.test_modeling_bark.BarkModelIntegrationTests
```

LGTM ! Let's wait for all the check to pass and merge then! Thanks for the great work here and all the iterations!

Thank you all again for your guidance and patience 🙏 much appreciated.