Constrained Beam Search [*With* Disjunctive Decoding]
Author: cwkeam
Pull request: yes

# ***Disjunctive*** Positive Constraint Decoding

@patrickvonplaten @Narsil 

A follow up to #15416, as a complete fix to #14081, as requested [here](https://github.com/huggingface/transformers/pull/15416#issuecomment-1031803308). On top of the current implementation of constrained beam search, it now allows for the "disjunctive" generation mentioned in the original issue and integrates a keyword argument for `model.generate()`.

It actually also fixes a logical error that was in the previous PR. It was a simple one-line fix found [here](https://github.com/huggingface/transformers/pull/15761/files#r811574539). I was treating the `beam_idx` within a batch as if it can be used as a `batch_beam_idx` across all the batches.

## 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.
- [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?

## Example Uses

These examples are tested in `tests/test_generation_utils.py`.


### Example 1: Simple Strong Constraint

Objective: Use `force_words_ids` keyword argument to force the word `"sind"` in the output.

Here `force_words_ids` is a `List[List[int]]`. This is in direct analogy to `bad_words_ids`. It's a much simpler use-case and I think that documentation & publicization of this feature would show this use to introduce users about it.

```python
@slow
def test_constrained_beam_search_example_translation_mixin(self):
    tokenizer = AutoTokenizer.from_pretrained("t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

    encoder_input_str = "translate English to German: How old are you?"
    force_words = ["sind"]

    input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
    force_words_ids = tokenizer(force_words, add_special_tokens=False, return_tensors="pt").input_ids

    outputs = model.generate(
        input_ids,
        force_words_ids=force_words_ids,
        num_beams=10,
        num_return_sequences=1,
        no_repeat_ngram_size=1,
        remove_invalid_values=True,
    )

    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    self.assertListEqual(outputs, ["Wie alter sind Sie?"])

```


### Example 2: A Mix of Strong Constraint and a Disjunctive Constraint.

Objective: complete the input sentence such that it must include the word `"scared"` and includes any of the word forms of `"scream"` (i.e. `["scream", "screams", "screaming", "screamed"]`)

Here `force_words_ids` is a `List[List[List[int]]]`. It's a list of disjunctive constraints, each of which is a list of words, and words are obviously a list of ids.
 

```python
def test_constrained_beam_search_mixed_mixin(self):
    model = GPT2LMHeadModel.from_pretrained("gpt2").to(torch_device)
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

    force_word = "scared"
    force_flexible = ["scream", "screams", "screaming", "screamed"]

    force_words_ids = [
        tokenizer([force_word], add_prefix_space=True, add_special_tokens=False).input_ids,
        tokenizer(force_flexible, add_prefix_space=True, add_special_tokens=False).input_ids,
    ]

    starting_text = ["The soldiers", "The child"]

    input_ids = tokenizer(starting_text, return_tensors="pt").input_ids.to(torch_device)

    outputs = model.generate(
        input_ids,
        force_words_ids=force_words_ids,
        num_beams=10,
        num_return_sequences=1,
        no_repeat_ngram_size=1,
        remove_invalid_values=True,
    )

    generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)

    self.assertListEqual(
        generated_text,
        [
            "The soldiers, who were all scared and screaming at each other as they tried to get out of the",
            "The child was taken to a local hospital where she screamed and scared for her life, police said.",
        ],
    )
```

## Future Improvements

I'm thinking of adding `OrderedConstraints`, `TemplateConstraints` further down the line. Currently, the generation is fulfilled by including the sequences, wherever in the output. For example, the above example has one sequence with scared -> screaming and the other has screamed -> scared. `OrderedConstraints` will allow the user to specify the order in which these constraints are fulfilled. 

`TemplateConstraints` will allow for a more niche use of the feature, where the objective can be something like:

```python
starting_text = "The woman"
template = ["the", "", "School of", "", "in"]

possible_outputs == [
   "The woman attended the Ross School of Business in Michigan.",
   "The woman was the administrator for the Harvard school of Business in MA."
]
```

or:
```python
starting_text = "The woman"
template = ["the", "", "", "University", "", "in"]

possible_outputs == [
   "The woman attended the Carnegie Mellon University in Pittsburgh.",
]
impossible_outputs == [
  "The woman attended the Harvard University in MA."
]
``` 

or if the user does not care about the number of tokens that can go in between two words, then one can just use `OrderedConstraint`.

But these are just minor things I might try adding, among the plethora of possible `Constraint` classes that anyone could add. The current implementation of `constrained_beam_search` is ***flexible*** enough to allow for any crazy way to constrain the beam search, as long as the `Constraint` object adheres to the ***strict*** definitions of `Constraint` as defined in the abstract base class.

Thanks for reviewing!






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

@patrickvonplaten Yes, I was worried about that part as well. It does seem kind of dangerous to allow two different behaviors given the input types to the same keyword argument.

But here were the reasons why I went with the current direction:
1. In general I feel like the approach that clutters the keyword args less to the already cluttered `generate()` function is better.
2. In disagreement with what you said above, I felt that it's not that perverse to call disjunctive constraints as something like a "word to be forced", since something like `["scream", "screamed", ...]` is supposed to be representing just one word after all, while trying to remain flexible to word forms.
3. I don't think it reduces much complexity (at least enough to justify hard-to-understand nascent keyword names & cluttering keyword args) to separate it into two keyword arguments. A user that understands the feature would use either systems equally effectively, and a user that hasn't bothered to learn about the disjunctive case could treat `force_words_ids` as a direct analogy to `bad_words_ids` and live without problems.

Though if the team agrees there should be separate keyword args then I'm always open to switch it up! I think more opinion might be needed.

> @patrickvonplaten Yes, I was worried about that part as well. It does seem kind of dangerous to allow two different behaviors given the input types to the same keyword argument.
> 
> But here were the reasons why I went with the current direction:
> 
> 1. In general I feel like the approach that clutters the keyword args less to the already cluttered `generate()` function is better.
> 2. In disagreement with what you said above, I felt that it's not that perverse to call disjunctive constraints as something like a "word to be forced", since something like `["scream", "screamed", ...]` is supposed to be representing just one word after all, while trying to remain flexible to word forms.
> 3. I don't think it reduces much complexity (at least enough to justify hard-to-understand nascent keyword names & cluttering keyword args) to separate it into two keyword arguments. A user that understands the feature would use either systems equally effectively, and a user that hasn't bothered to learn about the disjunctive case could treat `force_words_ids` as a direct analogy to `bad_words_ids` and live without problems.
> 
> Though if the team agrees there should be separate keyword args then I'm always open to switch it up! I think more opinion might be needed.

Good arguments! Ok to leave as it is then :-) 

From my side everything looks good. I agree with most of @Narsil's comments, but I also do think we can correct this in a follow-up PR. 

Happy to leave it up to you and @Narsil to decide on what to change before merging :-)

Hi @Narsil 

Don't worry about it! Actually I appreciate the clear-cutting comments. I'm really doing these PRs for personal learning purposes as a first-year college student and the feedback has been immensely helpful. Writing maintainable code is definitely something I'm trying to practice, so I really appreciate your thorough comments.

As for the type checking, from the very start of this PR, I've felt uncomfortable about how buggy it looks and the problems this could lead to. The reason why it hasn't been easy is because, in the generation process, the `token_ids` stored inside each `Constraint` object ends up interacting with `input_ids`, which is always a tensor. 

But I also couldn't make `token_ids` completely into tensors, because for the disjunctive case, it needs a list of list of words, whose varying sizes make it impossible to treat it as a 2D tensor. Also, in the user-side, it might be equally intuitive to put in a list of ints but also a tensor with a list of ids, and I wanted to cover that case too.

After giving it more thought with the feedback, I agree with you that the complexity is just not worth covering all the cases. I've overhauled all the code so that `Constraint` and everything related to is is ***completely torch free***. The constraints file does not even `import torch` anymore. All future implementations will need to follow this paradigm and, despite potential ineffiencies, it gives us maintanable consistency.



@Narsil thanks for the final check and thanks again for your thorough feedback!

@patrickvonplaten as soon as this merges, I’ll get started on a blog post! I just got it started by copying what you did for your [post on text generation](https://huggingface.co/blog/how-to-generate). Is writing everything with a similar format in a colab notebook the pushing it to huggingface/blogs repo the right way to proceed?

Softly pinging @patrickvonplaten for merge to really get going with the blog & future additions to this feature. Also wondering if rebasing with master again is necessary since its been a week since.

Thanks for the ping - thought this was already merged. Looked over it again and everything looks good to me :-) @Narsil also approved the PR so merging now. Great job!

Does this PR support the application of different constraints for different elements in a batch? For example, `starting_text = ["The soldiers", "The child"]` where I want to force the word `scared` only for the first item and `played` only for the second item.

I am trying to apply the constrained beam search method on a large dataset, with different constraints for each instance. Using a batch size of 1 seems to be the only solution as of now.