在🤗 Transformers中使用受限束搜索引导文本生成

在🤗 Transformers中使用受限束搜索来引导文本生成

介绍

本博文假设读者熟悉使用不同变种的束搜索进行文本生成的方法,如博文中所解释的:“如何生成文本:使用变种解码方法进行语言生成与Transformer”

与普通束搜索不同,约束束搜索允许我们对文本生成的输出施加控制。这很有用,因为有时我们对输出内部的内容非常明确。例如,在神经机器翻译任务中,我们可能通过字典查找知道最终翻译中必须包含哪些单词。有时,对于语言模型来说,几乎同样可能的生成输出对于最终用户来说可能并不同样理想,因为具体的上下文。通过允许用户告诉模型最终输出中必须包含哪些单词,可以解决这两种情况。

为什么很困难

然而,这实际上是一个非常棘手的问题。这是因为该任务要求我们在生成的最终输出的某个位置、某个时刻强制生成某些子序列。

假设我们想生成一个包含短语 p 1 = { t 1 , t 2 } p_1=\{ t_1, t_2 \} p 1 ​ = { t 1 ​ , t 2 ​ } 的句子 S 。让我们将期望的句子 S S S 定义为:

S e x p e c t e d = { s 1 , s 2 , . . . , s k , t 1 , t 2 , s k + 1 , . . . , s n } S_{expected} = \{ s_1, s_2, …, s_k, t_1, t_2, s_{k+1}, …, s_n \} S e x p e c t e d ​ = { s 1 ​ , s 2 ​ , . . . , s k ​ , t 1 ​ , t 2 ​ , s k + 1 ​ , . . . , s n ​ }

问题在于,束搜索逐个标记地生成序列。尽管不完全准确,但可以将束搜索视为函数 B ( s 0 : i ) = s i + 1 B(\mathbf{s}_{0:i}) = s_{i+1} B ( s 0 : i ​ ) = s i + 1 ​ ,其中它从 0 0 0 到 i i i 查看当前生成的标记序列,然后预测 i + 1 i+1 i + 1 处的下一个标记。但是在任意的步骤 i < k i < k i < k ,这个函数怎么知道这些标记必须在一些未来的步骤 k k k 处生成?或者当它在步骤 i = k i=k i = k 时,它怎么能确定这是强制生成标记的最佳位置,而不是某个未来的步骤 i > k i>k i > k ?

如果有多个要求不同的约束怎么办?如果您想强制生成短语 p 1 = { t 1 , t 2 } p_1=\{t_1, t_2\} p 1 ​ = { t 1 ​ , t 2 ​ } ,还想强制生成短语 p 2 = { t 3 , t 4 , t 5 , t 6 } p_2=\{ t_3, t_4, t_5, t_6\} p 2 ​ = { t 3 ​ , t 4 ​ , t 5 ​ , t 6 ​ } 该怎么办?如果您希望模型在这两个短语之间选择怎么办?如果我们想强制短语 p 1 p_1 p 1 ​ ,并在短语列表 { p 21 , p 22 , p 23 } \{p_{21}, p_{22}, p_{23}\} { p 2 1 ​ , p 2 2 ​ , p 2 3 ​ } 中只强制生成一个短语怎么办?

上述示例实际上是非常合理的用例,如下所示,新的约束束搜索功能可以实现所有这些!

本文将快速介绍新的约束束搜索功能对您的作用,然后深入介绍它在内部是如何工作的。

示例1:强制使用单词

假设我们正在尝试将"How old are you?"翻译成德语。

在非正式场合下,您会说"Wie alt bist du?",而在正式场合下,您会说"Wie alt sind Sie?"

根据上下文,我们可能希望一种形式的礼貌而不是另一种形式,但是我们如何告诉模型呢?

下面是我们在传统束搜索设置中进行文本翻译的方式。

!pip install -q git+https://github.com/huggingface/transformers.git

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids

outputs = model.generate(
    input_ids,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)

print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt bist du?

但是,如果我们知道我们想要正式的输出而不是非正式的输出怎么办?如果我们从先前的知识中知道生成必须包含的内容,并且我们可以将其注入到生成中怎么办?

现在,通过model.generate()force_words_ids关键字参数,我们可以实现以下功能:

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 = ["Sie"]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids
force_words_ids = tokenizer(force_words, add_special_tokens=False).input_ids

outputs = model.generate(
    input_ids,
    force_words_ids=force_words_ids,
    num_beams=5,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

如您所见,我们能够通过先前对所需输出的知识进行指导生成。以前,我们必须生成一堆可能的输出,然后筛选符合要求的输出。现在,我们可以在生成阶段做到这一点。

示例2:分离约束

我们上面提到了一种用例,即我们知道我们想要在最终输出中包含哪些单词。其中一个例子可能是在神经机器翻译期间使用词典查找。

但是,如果我们不知道要使用哪些词形,在这种情况下,我们希望像["raining", "rained", "rains", ...]这样的输出同样可能?从更一般的意义上说,总会有一些情况,我们不希望精确逐字逐句地使用确切的单词,而是可能接受其他相关的可能性。

允许此行为的约束是分离约束,它允许用户输入一个单词列表,其目的是指导生成,使得最终输出必须至少包含列表中的一个单词。

下面是一个使用上述两种类型约束的示例:

from transformers import GPT2LMHeadModel, GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained("gpt2")
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

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


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(tokenizer.decode(outputs[1], skip_special_tokens=True))

Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.


Output:
----------------------------------------------------------------------------------------------------
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.

正如您所看到的,第一个输出使用了"screaming",第二个输出使用了"screamed",而且两者都直接使用了"scared"。要选择的列表["screaming", "screamed", ...]不必是单词形式,这可以满足任何需要从单词列表中选择一个单词的用例。

传统 Beam 搜索

以下是传统Beam 搜索的示例,摘自之前的博客文章:

与贪婪搜索不同,Beam 搜索通过保留一个更长的假设列表来工作。在上面的图片中,我们在生成的每个可能步骤中显示了三个可能的下一个标记。

下面是上面示例中 Beam 搜索的第一步的另一种观察方式,假设num_beams=3

贪婪搜索只会选择"The dog",而 Beam 搜索会进一步考虑"The nice""The car"

在下一步中,我们考虑前一步中创建的三个分支的下一个可能的标记。

尽管我们最终考虑的输出数量比num_beams要多得多,但在该步骤结束时,我们将其减少到num_beams。我们不能一直扩展下去,否则我们需要跟踪的beams数量将会迅速变得非常庞大(在10步之后,10个beams变为10,000,000,000个beams!)。

对于剩余的生成部分,我们重复上述步骤,直到满足结束条件,例如生成<eos>标记或达到max_length。扩展、排序、缩减和重复。

约束 Beam 搜索试图通过在生成的每一步中注入所需的标记来满足约束条件。

假设我们试图强制生成输出中的短语"is fast"

在传统的 Beam 搜索设置中,我们在每个分支中找到前k个最有可能的下一个标记,并将它们附加以供考虑。在约束设置中,我们同样这样做,但还将接近满足约束条件的标记附加上去。以下是一个演示:

除了像"dog""nice"这样的通常概率较高的下一个标记外,我们还强制使用"is"这个标记,以便更接近满足我们的约束条件"is fast"

对于下一步,下面的分支候选者与传统 Beam 搜索中的大部分情况相同。但是与上面的示例一样,约束 Beam 搜索通过在每个新的分支中强制约束条件,添加到现有的候选者中:

银行

在讨论下一步之前,我们需要考虑上述步骤中出现的不希望的行为。

仅仅简单地强制所期望的短语"is fast"出现在输出中的问题是,大多数情况下,您最终会得到像上面的"The is fast"这样毫无意义的输出。这实际上使得解决这个问题变得不平凡。关于解决这个问题的复杂性的更深入讨论可以在最初提出的功能请求问题中找到,该问题在huggingface/transformers中提出。

银行通过在满足约束条件和生成合理输出之间取得平衡来解决这个问题。

银行 n n n 指的是在满足约束条件方面取得 n n n 步进度的梁的列表。在将所有可能的梁排序并放入各自的银行后,我们进行循环选择。在上面的示例中,我们会从银行 2 中选择最有可能的输出,然后从银行 1 中选择最有可能的,再从银行 0 中选择一个,从银行 2 中选择第二个最有可能的,从银行 1 中选择第二个最有可能的,以此类推。由于我们使用了 num_beams=3 ,我们只需要重复上述过程三次,就可以得到 ["The is fast", "The dog is", "The dog and"]

通过这种方式,即使我们强制模型考虑我们手动添加所需标记的分支,我们仍然可以跟踪其他高概率序列,这些序列可能更有意义。即使 "The is fast" 完全满足我们的约束条件,但它并不是一个非常合理的短语。幸运的是,我们在将来的步骤中可以使用 "The dog is""The dog and" ,希望这些短语能在以后产生更合理的输出。

这种行为在上述示例的第三步中得到了证明:

请注意,"The is fast" 不需要手动添加任何约束标记,因为它已经满足了(即已经包含了短语 "is fast" )。还请注意,像 "The dog is slow" 或者 "The dog is mad" 这样的梁实际上在银行 0 中,因为虽然它包含了标记 "is" ,但它必须从头开始生成 "is fast" 。通过在 "is" 后面添加类似于 "slow" 的东西,它实际上重置了自己的进度。

最后请注意,我们最终得到的是包含我们约束短语的合理输出:"The dog is fast"

最初我们担心盲目添加所需的标记会导致类似于 "The is fast" 的无意义短语。然而,通过从银行进行循环选择,我们隐式地摆脱了无意义的输出,而更倾向于更有意义的输出。

更多关于约束类和自定义约束的信息

从上述解释中可以得出以下要点。在每一步中,我们都不断要求模型考虑满足我们约束的标记,同时跟踪不满足约束的梁,直到我们得到包含所需短语的相对高概率序列。

因此,设计这个实现的一种有原则的方法是将每个约束表示为一个 Constraint 对象,其目的是跟踪其进度并告诉梁搜索要生成的下一个标记。尽管我们为 model.generate() 提供了关键字参数 force_words_ids ,但实际上发生的是:

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, PhrasalConstraint

tokenizer = AutoTokenizer.from_pretrained("t5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")

encoder_input_str = "translate English to German: How old are you?"

constraints = [
    PhrasalConstraint(
        tokenizer("Sie", add_special_tokens=False).input_ids
    )
]

input_ids = tokenizer(encoder_input_str, return_tensors="pt").input_ids


outputs = model.generate(
    input_ids,
    constraints=constraints,
    num_beams=10,
    num_return_sequences=1,
    no_repeat_ngram_size=1,
    remove_invalid_values=True,
)


print("Output:\n" + 100 * '-')
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Output:
----------------------------------------------------------------------------------------------------
Wie alt sind Sie?

您可以自己定义一个并将其输入到 constraints 关键字参数中,以设计您的独特约束。您只需要创建一个 Constraint 抽象接口类的子类,并遵循其要求。您可以在此处找到有关 Constraint 的定义的更多信息。

一些独特的想法(尚未实现; 也许你可以试试!)包括像OrderedConstraintsTemplateConstraints这样的约束,可能会在以后添加。目前,生成是通过包含序列来实现的,无论在输出中的任何位置。例如,先前的示例中有一个序列从“scared”到“screaming”,另一个序列从“screamed”到“scared”。OrderedConstraints可以允许用户指定这些约束的实现顺序。

TemplateConstraints可以允许更特定的功能使用,例如:

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."
]

或者:

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."
]

或者,如果用户不关心两个词之间可以插入多少个令牌,那么可以使用OrderedConstraint

结论

有约束的波束搜索为我们注入外部知识和要求到文本生成提供了灵活的方式。以前,没有简单的方法告诉模型:1.包括一系列序列;2.其中一些是可选的,一些是必需的;3.它们在序列中的合理位置生成。现在,我们可以通过不同子类的Constraint对象完全控制生成过程!

这个新功能主要基于以下论文:

  • Guided Open Vocabulary Image Captioning with Constrained Beam Search
  • Fast Lexically Constrained Decoding with Dynamic Beam Allocation for Neural Machine Translation
  • Improved Lexically Constrained Decoding for Translation and Monolingual Rewriting
  • Guided Generation of Cause and Effect

像上面这些一样,许多新的研究论文正在探索使用外部知识(例如,知识图谱,知识库)来指导大型深度学习模型的输出的方法。希望这个有约束的波束搜索功能能成为另一种实现这一目的的有效方法。

感谢所有为这个功能做出指导的人:Patrick von Platen从最初的问题到最终的PR都参与其中,Narsil Patry为代码提供了详细的反馈。

本文的缩略图使用了一个包含归属的图标:Shorthand icons created by Freepik – Flaticon