Python 如何防止transformer generate函数产生某些字

发布于03月28日

我有以下code条:

from transformers import T5Tokenizer, T5ForConditionalGeneration

tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")

input_ids = tokenizer("The <extra_id_0> walks in <extra_id_1> park", return_tensors="pt").input_ids

sequence_ids = model.generate(input_ids)
sequences = tokenizer.batch_decode(sequence_ids)
sequences

目前，它生产的产品是:

['<pad><extra_id_0> park offers<extra_id_1> the<extra_id_2> park.</s>']

有没有办法防止生成器产生某些单词(例如stopwords = ["park", "offer"])？

推荐答案

在查看文档后发现，您可以在generate()中传递一个bad_words_ids参数

给定一个坏单词列表，您可以使用以下命令创建id列表

tokenizer(bad_words, add_special_tokens=False).input_ids

input_ids = tokenizer("The <extra_id_0> walks in <extra_id_1> park", return_tensors="pt").input_ids
bad_words = ["park", "offers"]
bad_words_ids = tokenizer(bad_words, add_special_tokens=False).input_ids 
#[[2447], [704]]

sequence_ids = model.generate(input_ids, bad_words_ids=bad_words_ids)
#tensor([[    0, 32099,  1061,    19,     3,     9,   710,  1482,   550,    45, 32098,     8, 32097,  1061,     5,     1]])

sequences = tokenizer.batch_decode(sequence_ids)
print(sequences) 
#['<pad><extra_id_0> Park is a short walk away from<extra_id_1> the<extra_id_2> Park.</s>']

Colab demo个