Constrained Beam Search [without disjunctive decoding]

[cwkeam]

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).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. [Feature Contribution] Disjunctive Positive Constraint Decoding (adding force_tokens to model.generate()) #14081
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• 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:

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:

# 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)
...

[HuggingFaceDocBuilder]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

Very cool taking a look now

[patrickvonplaten]

It's not super obvious to me what this contraint would do actually. Could we maybe use a "real" contraint here?

[cwkeam]

What do you mean by a "real" constraint as opposed to what is in the above example?

I do share the general concern about the lack of clarity on what specific effect these "constraints" have on the generation output. As per the other comment below about showing the actual output. I'll rework the documentation again in a following commit.

[cwkeam]

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!

[patrickvonplaten]

This code is very sensitive and should not be changed at all. Is GitHub just showing a very bad diff or was some code changed here?

[patrickvonplaten]

Looks good to me. IMO we can merge this.

@Narsil , @patil-suraj @sgugger @LysandreJik - could you take a very high level look as well?

[patrickvonplaten]

Ok actually this looks good. Just seems like GitHub messed up the diff, no?

[patrickvonplaten]

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 :-)

[patrickvonplaten]

See https://github.com/huggingface/transformers/pull/15416/files#r800972360 - after than we can merge I think

[patrickvonplaten]

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

[sgugger]

Thanks a lot for your PR! It's great that you put so many comments in your code, but we use multi-line comments inside the code, not triple-quoted strings. Those will be confused with docstrings by our utils, so you should switch them back to comments.

[cwkeam]

@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]

@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]

Awesome - thanks a lot for reverting this change and opening an issue instead.
As a final step, could you maybe try to apply @sgugger's comments about docstrings and comments? :-) Happy to help / correct it if you have questions

[cwkeam]

@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

[patrickvonplaten]

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 🚀