# Example of prefix\_allowed\_tokens\_fn() while text generation

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fexample-of-prefix-allowed-tokens-fn-while-text-generation%2F6635%3E%5C
**Category:** 🤗Transformers\
**Created:** [June 7, 2021, 9:06pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fexample-of-prefix-allowed-tokens-fn-while-text-generation%2F6635 "2021-06-07T21:06:57Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![vnik18](https://avatars.discourse-cdn.com/v4/letter/v/6a8cbe/32.png) [@vnik18](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fvnik18)%5C
**Post date:** [June 7, 2021, 9:06pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fexample-of-prefix-allowed-tokens-fn-while-text-generation%2F6635%2F1 "2021-06-07T21:06:57Z")

</div>

Hello, I would like to use the ` prefix_allowed_tokens_fn` as an input to the `model.generate()` function, in order to perform constrained text generation with BART.

> <https://github.com/huggingface/transformers/blob/master/src/transformers/generation_utils.py#L670>

I tried to adapt the function in the original repository [here](https://github.com/facebookresearch/GENRE/tree/main/examples_genre), but it doesn’t seem to be working. Can you please tell me if there are any examples of the kinds of functions that can be given as input to this parameter? Thank you!

---

<div class="post-metadata">

**Author:** ![keshavkolluru](https://avatars.discourse-cdn.com/v4/letter/k/f9ae1b/32.png) [@keshavkolluru](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fkeshavkolluru)%5C
**Post date:** [July 18, 2022, 5:45am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fexample-of-prefix-allowed-tokens-fn-while-text-generation%2F6635%2F2 "2022-07-18T05:45:16Z")

</div>

Probably too late but just in case it helps someone else - the following code worked for me:

> from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer  
> from genre.trie import MarisaTrie
> 
> model = AutoModelForSeq2SeqLM.from\_pretrained(‘t5-base’)  
> tokenizer = AutoTokenizer.from\_pretrained(‘t5-base’)
> 
> trie = MarisaTrie([[0]+tokenizer.encode(‘Hello World’)])
> 
> output = model.generate(tokenizer.encode(‘Hello World’, return\_tensors=‘pt’), prefix\_allowed\_tokens\_fn=lambda batch\_id, sent: trie.get(sent.tolist()))

The above snipped will always produce “Hello World” as the output. You can also include multiple strings when creating the Marisa trie.

The [0] is required at the start as the t5 model always produces 0 as the first token.

The definition of the trie is taken from here: [https://github.com/facebookresearch/GENRE/blob/main/genre/trie.py](https://github.com/facebookresearch/GENRE/blob/main/genre/trie.py) and the file requires ‘pip install marisa-trie’ to be installed in the environment.

---

<div class="post-metadata">

**Author:** ![bryanzhou008](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bryanzhou008/32/5463_2.png) [@bryanzhou008](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fbryanzhou008)%5C
**Post date:** [July 21, 2022, 8:33pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fexample-of-prefix-allowed-tokens-fn-while-text-generation%2F6635%2F3 "2022-07-21T20:33:02Z")

</div>

Thank you so much for this answer, it is exactly what I‘ve been searching for!