Constrained Beam Search [without disjunctive decoding]
Author: cwkeam
Pull request: yes

# ~~Disjunctive~~ Positive Constraint Decoding 

@patrickvonplaten @Narsil 

Fixes #14081, 

But without the "disjunctive" component. I think it's a good idea to first  integrate this `Constraint` framework before adding more advanced techniques like Disjunctive Decoding, which can arguably be added with ease afterwards.

Currently this PR is able to force the generation of tokens and/or an ordered sequence of tokens in the beam search with `constrained_beam_search` in `GenerationMixin.generate`. It finds the best possible, highest probable list of beam outputs in fulfilling a list of `Constraint` objects, yielding results more sensible than just blindly forcing the tokens in the output sequence.

## 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/master/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. #14081 
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/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.

Tagging @patrickvonplaten because my main back-and-forth was with you on this PR.

## Example Use

Here is an example of how one could use this functionality:

```python
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from transformers.generation_beam_constraints import PhrasalConstraint

device = "cuda"

model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

force_text = " big monsters"
force_text_2 = " crazy"
force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device)[0]
force_tokens_2 = tokenizer.encode(force_text_2, return_tensors="pt").to(device)[0]

constraints = [
    PhrasalConstraint(force_tokens),
    PhrasalConstraint(force_tokens_2)
]

input_text = ["The baby is crying because"] 

model_inputs = tokenizer(input_text, return_tensors="pt")

for key, value in model_inputs.items():
    model_inputs[key] = value.to(device)

k = model.generate(
    **model_inputs,
    constraints=constraints,
    num_beams=7,
    num_return_sequences=7,
    no_repeat_ngram_size=2
)

for out in k:
    print(tokenizer.decode(out))
```
For some example outputs:
```
The baby is crying because she's been told crazy big monsters are going to come and kill her.
The baby is crying because she's been told crazy big monsters are coming for her.
The baby is crying because she's been told crazy big monsters are going to come after her.
```

## The Constraint Framework

Users can define their own constraints by inheriting the `Constraint` abstract class and this framework is ensured to work as desired, because the `Constraint` class is quite strictly defined. If an implementation passes the `self.test()` function of this interface then it necessarily works as desired. An incorrect implementation will lead to an error. But probably most users would be fine just using the `TokenConstraint` and `PhrasalContraint` classes.

## Testing

My approach to testing this new functionality was that it passes every test that was related to the existing `beam_search` function, with the **additional testing that requires that the generated output indeed fulfills the constraint.**

One such example:
```python
# tests/test_generation_utils.py 

def test_constrained_beam_search_generate(self):
    for model_class in self.all_generative_model_classes:
        ....
        # check `generate()` and `constrained_beam_search()` are equal for `num_return_sequences`
        # Sample constraints
        force_tokens = torch.randint(min_id, max_id, (1, 2)).type(torch.LongTensor)[0]
        constraints = [
            PhrasalConstraint(force_tokens),
        ]
        
        num_return_sequences = 2
        max_length = 10
        
        beam_kwargs, beam_scorer = self._get_constrained_beam_scorer_and_kwargs(
            input_ids.shape[0], max_length, constraints, num_return_sequences=num_return_sequences
        )
        
        output_generate, output_beam_search = self._constrained_beam_search_generate(
            model=model,
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=max_length,
            constrained_beam_scorer=beam_scorer,
            constraints=constraints,
            beam_kwargs=beam_kwargs,
            logits_processor=logits_processor,
            logits_process_kwargs=logits_process_kwargs,
        )
        self.assertListEqual(output_generate.tolist(), output_beam_search.tolist())
        
        for generation_output in output_generate:
            self._check_sequence_inside_sequence(force_tokens, generation_output)
...
```




Comments:
_The documentation is not available anymore as the PR was closed or merged._

Very cool taking a look now

Hey @cwkeam,

This already looks like a super clean and impactful addition to the `transformers` library! Great job. I really like how you implement this new generation method in a extendable way without affecting existing generate methods. I'm sure we can merge it soon. 

At the moment, it seems like your example above doesn't work. When I try to run:

```python
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from transformers.generation_beam_constraints import PhrasalConstraint

device = "cuda"

model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

force_text = " big monsters"
force_text_2 = " crazy"
force_tokens = tokenizer.encode(force_text, return_tensors="pt").to(device)[0]
force_tokens_2 = tokenizer.encode(force_text_2, return_tensors="pt").to(device)[0]

constraints = [
    PhrasalConstraint(force_tokens),
    PhrasalConstraint(force_tokens_2)
]

input_text = ["The baby is crying because"] 

model_inputs = tokenizer(input_text, return_tensors="pt")

for key, value in model_inputs.items():
    model_inputs[key] = value.to(device)

k = model.generate(
    **model_inputs,
    constraints=constraints,
    num_beams=7,
    num_return_sequences=7,
    no_repeat_ngram_size=2
)

for out in k:
    print(tokenizer.decode(out))
```

on my computer, I'm getting the following error:

```
ValueError: `token_ids` has to be list of positive integers or a `torch.LongTensor` but is tensor([7165], device='cuda:0')For single token constraints, refer to `TokenConstraint`.
```

I think the example is really nice and exactly what we are looking for. So maybe you could also add it as a `@slow` test under `tests/test_generation_utils.py` once it works :-)

Let me know if you need further help!

Added 1. clear examples (**and tests for them**, 2. more specific input parameter tests, and 3. **bug fixes in other beam search code**.

## 1. Examples (text completion & translation)

The latter example really emphasizes the usefulness of this feature.

Fixed up the first example as requested previously, showing its ability that it can

- [x] handle multiple constraints
- [x] handle constraints that have more than one words (sequence of tokens)

```python
def test_constrained_beam_search(self):
    model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device)
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

    force_tokens = tokenizer.encode(" scared", return_tensors="pt").to(torch_device)[0]
    force_tokens_2 = tokenizer.encode(" big weapons", return_tensors="pt").to(torch_device)[0]

    constraints = [
        PhrasalConstraint(force_tokens),
        PhrasalConstraint(force_tokens_2),
    ]

    starting_text = ["The soldiers were not prepared and"]

    input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)

    outputs = model.generate(
        input_ids,
        constraints=constraints,
        num_beams=10,
        num_return_sequences=1,
        no_repeat_ngram_size=1,
        max_length=30,
        remove_invalid_values=True,
    )

    generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    self.assertListEqual(
        generated_text,
        [
            "The soldiers were not prepared and didn't know how big the big weapons would be, so they scared them off. They had no idea what to do",
        ],
    )

```

and here's a translation example:

```python
def test_constrained_beam_search_example_integration(self):
    tokenizer = AutoTokenizer.from_pretrained("t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

    encoder_input_str = "translate English to German: How old are you?"
    encoder_input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

    # lets run beam search using 5 beams
    num_beams = 5
    # define decoder start token ids
    input_ids = torch.ones((num_beams, 1), device=model.device, dtype=torch.long)
    input_ids = input_ids * model.config.decoder_start_token_id

    # add encoder_outputs to model keyword arguments
    model_kwargs = {
        "encoder_outputs": model.get_encoder()(
            encoder_input_ids.repeat_interleave(num_beams, dim=0), return_dict=True
        )
    }

    constraint_str = "sind"
    constraint_token_ids = tokenizer.encode(constraint_str)[:-1]  # remove eos token
    constraints = [PhrasalConstraint(token_ids=constraint_token_ids)]

    # instantiate beam scorer
    beam_scorer = ConstrainedBeamSearchScorer(
        batch_size=1, num_beams=num_beams, device=model.device, constraints=constraints
    )

    # instantiate logits processors
    logits_processor = LogitsProcessorList(
        [
            MinLengthLogitsProcessor(5, eos_token_id=model.config.eos_token_id),
        ]
    )

    outputs = model.constrained_beam_search(
        input_ids, beam_scorer, constraints=constraints, logits_processor=logits_processor, **model_kwargs
    )
    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    self.assertListEqual(outputs, ["Wie alter sind Sie?"])

```

This is a really useful example, because the following is the test for **ordinary beam search**

```python
def test_beam_search_example_integration(self):
    # ... the exact same code as above minus constraints
    outputs = model.beam_search(input_ids, beam_scorer, logits_processor=logits_processor, **model_kwargs)
    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    self.assertListEqual(outputs, ["Wie alt bist du?"])

```

I've never studied German before, but upon looking it up, it seemed like there were two equally realistic ways to translate "How old are you" in German. While the ordinary beam search output may be fine for most uses, maybe the user specifically wants the other way of translating it. In this case, they can do as I have and force the generation to include the token that changes the direction of the translation as desired.

## 2. Input Parameters tests

Added @patrickvonplaten 's recommendation as well as dealing with `num_groups`:

```python
if num_beams <= 1:
    raise ValueError("`num_beams` needs to be greater than 1 for constrained genertation.")

if do_sample:
    raise ValueError("`do_sample` needs to be false for constrained generation.")

if num_beam_groups is not None and num_beam_groups > 1:
    raise ValueError("`num_beam_groups` not supported yet for constrained generation.")
```

I'm not sure whether the last part is a good idea and whether I should just make the implementation work for `num_beam_groups>1` but the fact is that it can't handle that case yet (will throw an error).

## 3. Bug Fixes in **beam_search** with `max_length`

As I was writing the test code for the few examples above, I realized that the example in the [huggingface documentation](https://huggingface.co/docs/transformers/v4.16.2/en/main_classes/model#transformers.generation_utils.GenerationMixin.beam_search) for **ordinary beam search** was actually failing, because of the **recent change with `max_length`**.

The above testing code `test_beam_search_example_integration(self)` is taken directly from the huggingface documentation for `GenerationMixin.beam_search` and it fails, because `max_length` can be set to None unless otherwise specified, and it throws an error saying that you can't compare a `Tensor` to a `NoneType`.

[From `BeamSearchScorer.finalize()`](https://github.com/huggingface/transformers/blob/master/src/transformers/generation_beam_search.py#L333),

```python
sent_max_len = min(sent_lengths.max().item() + 1, max_length) # ERROR 1
decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)
if sent_lengths.min().item() != sent_lengths.max().item():
    assert pad_token_id is not None, "`pad_token_id` has to be defined"
    decoded.fill_(pad_token_id)

for i, hypo in enumerate(best):
    decoded[i, : sent_lengths[i]] = hypo
    if sent_lengths[i] < max_length: # ERROR 2
        decoded[i, sent_lengths[i]] = eos_token_id
```

So I changed that to:

```python
sent_lengths_max = sent_lengths.max().item() + 1
sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max
decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len)
# shorter batches are padded if needed
if sent_lengths.min().item() != sent_lengths.max().item():
    assert pad_token_id is not None, "`pad_token_id` has to be defined"
    decoded.fill_(pad_token_id)

# fill with hypotheses and eos_token_id if the latter fits in
for i, hypo in enumerate(best):
    decoded[i, : sent_lengths[i]] = hypo
    if sent_lengths[i] < sent_max_len:
        decoded[i, sent_lengths[i]] = eos_token_id
```







Hey @patrickvonplaten 

I also wanted to ask your opinions on making this feature integrate more easily to the `model.generate()` function. The same way how `bad_word_ids` is converted into a `LogitsProcessor` in the internal logic, maybe there's a way to add a similarly simple API like `force_token_ids` that internally converts it into `Constraint` objects in the backend.

I'm thinking of something like:

```python
starting_text = ["The soldiers were not prepared and"]
force_tokens = ["scared", "big weapons"]

input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)
force_token_ids = tokenizer(force_tokens).input_ids

outputs = model.generate(
    input_ids,
    force_token_ids=force_token_ids, # [[28349], [23981, 3721]] something that looks like this (made the numbers up)
   ...
)
```

But the problem arises when we start talking about **Disjunctive Positive Constraint Decoding** which was the main subject of the original #14081 issue. Although the disjunctive constraint generation implementation isn't even made yet, I think the best approach is something like this:

```python
from transformers.generation_beam_constraints import DisjunctiveConstraint

variations_1 = ["rain", "raining", "rained"]
variations_2 = ["smile", "smiling", "smiled"]

force_rain = tokenizer.encode(variations, return_tensors="pt")
force_smile = tokenizer.encode(variatiosn_2, return_tensors="pt")

constraints = [
    DisjunctiveConstraint(force_rain),
    DisjunctiveConstraint(force_smile)
]

outputs = model.generate(
    input_ids,
    constraints=constraints,
   ...
)
```
Where `DisjunctiveConstraint` is a constraint that is fulfilled by fulfilling just one among the list.

However, this feature would be something that's quite awkward to make into the `force_token_ids` API.

I think one way to approaching this is realizing that a `PhrasalConstraint` is basically a subset of a `DisjunctiveConstraint`:

```python
variations_1 = ["rain"]
force_rain = tokenizer.encode(variations, return_tensors="pt")

constraints = [
    DisjunctiveConstraint(force_rain),
]
```
The above has the same effect as just doing `PhrasalConstraint(tokenizer_encode("rain"))`, since the Disjunctive code would be choosing one from a list with one item.

So maybe the API above should be changed to:

```python
starting_text = ["The soldiers were not prepared and"]

force_tokens = [ 
    ["scared", "another"], 
    ["big weapons"]
]

input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)
force_token_ids = tokenizer(force_tokens).input_ids

outputs = model.generate(
    input_ids,
    force_token_ids=force_token_ids, # [
        [ [28349], [456849] ], # DisjunctiveConstraint, choose one between two
        [ [23981, 3721] ] # same ordinary PhrasalConstraint for a sequence, yet still a Disjunctive object.
     ]  
   ...
)
```

Then the input type for `force_token_ids` would be `List[List[List[int]]]` that sounds quite messy.

Let me know your thoughts on this!

Hey @patrickvonplaten,

That's a good point. I think that the `DisjunctiveConstraint` feature should also come as a separate PR for maintanence purposes as you said. I'll open a new PR soon.

The checks seem to be all passing now! Will you do a final review?

Thanks!



Incredible job at implementing this @cwkeam - it's very clean and enables some really cool use cases. 
If you'd like we could now the Conjunctive decoding in a follow-up PR then and make it a bit easier for the user to enable it via some keywords for `generate()` and then post about this method - I'm sure the community will be very excited about this addition. If you'd like to promote the integration already at this point, we could also do it right away and delay the conjuntive decoding PR a bit - as you wish :-)

All slow tests run correctly and the fast tests are not slowed down by this PR. One final change @cwkeam and I think then we are good to go

@patrickvonplaten @sgugger Thanks for the review! I've adjusted all the code accordingly.

As for promoting these changes, I think I'll save it for the `DisjunctiveConstraint` and simple integration to the `generate()` function is complete. I'll open a new PR about it soon. In terms of after those other changes are done, what would be a nice way to promote this change? Could it be possible to write a blog post in huggingface blogs as well as posting on TowardsDataScience about it?

Thanks so much for everyone's help!

> @patrickvonplaten @sgugger Thanks for the review! I've adjusted all the code accordingly.
> 
> As for promoting these changes, I think I'll save it for the `DisjunctiveConstraint` and simple integration to the `generate()` function is complete. I'll open a new PR about it soon. In terms of after those other changes are done, what would be a nice way to promote this change? Could it be possible to write a blog post in huggingface blogs as well as posting on TowardsDataScience about it?
> 
> Thanks so much for everyone's help!

A blog post would be a great idea indeed! This is a really nice feature to write a blog post about I think. Happy to help you 

@patrickvonplaten 

I've made all the requested adjustments on the comments. 

The test is failing right now at just one script:

```
tests/test_modeling_speech_encoder_decoder.py::Wav2Vec2Speech2Text2::test_encoder_decoder_model_generate
```

which actually runs fine on my machine. Not sure why it's throwing this error after pushing a new commit with adjustments to nothing but comments. This has happened before but simply worked out after fetching from master. How does one generally deal with these problems?

EDIT: upon pushing again with some docstrings adjustments, supposedly just by just triggering a re-run, it's passing now

Incredible work here @cwkeam - this was one of the cleanest and highest quality PRs that I've seen in a long time! Merging - let's try to tackle disjunctive decoding next :rocket: