Example of prefix_allowed_tokens_fn() while text generation
Hello, I would like to use theprefix_allowed_tokens_fnas an input to themodel.generate()function, in order to perform constrained text generation with BART.github.comhuggingface/transformers/blob/master/src/transformers/generation_utils.py#L670length_penalty: Optional[float] = None,no_repeat_ngram_size: Optional[int] = None,encoder_no_repeat_ngram_size: Optional[int] = None,num_return_sequences: Optional[int] = None,max_time: Optional[float] = None,max_new_tokens: Optional[int] = None,decoder_start_token_id: Optional[int] = None,use_cache: Optional[bool] = None,num_beam_groups: Optional[int] = None,diversity_penalty: Optional[float] = None,prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,output_attentions: Optional[bool] = None,output_hidden_states: Optional[bool] = None,output_scores: Optional[bool] = None,return_dict_in_generate: Optional[bool] = None,forced_bos_token_id: Optional[int] = None,forced_eos_token_id: Optional[int] = None,remove_invalid_values: Optional[bool] = None,synced_gpus: Optional[bool] = None,**model_kwargs,) -> Union[GreedySearchOutput, SampleOutput, BeamSearchOutput, BeamSampleOutput, torch.LongTensor]:I tried to adapt the function in the original repositoryhere, 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!
2022-07-18T05:45:16Z
Probably too late but just in case it helps someone else - the following code worked for me:from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizerfrom genre.trie import MarisaTriemodel = 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.pyand the file requires ‘pip install marisa-trie’ to be installed in the environment.
2022-07-21T20:33:02Z
Thank you so much for this answer, it is exactly what I‘ve been searching for!