从可能的选项中选择句子中缺少的单词

Question

我有一个缺少单词的句子。

"I want to buy a ___ to drive it."

我有答案选项可能在这里。

["car", "cat", "can"]

我需要选择最合适的词。我想用BERT来做，但是没有办法指定哪些单词可以。

Answer 1

您可以将 BERT 限制为单词列表。事实上，我做了一个库来做这个：FitBert。

from fitbert import FitBert

fb = FitBert()
masked_string = "I want to buy a ***mask*** to drive it."
options = ["car", "cat", "can"]
ranked_options = fb.rank(masked_string, options=options)
assert ranked_options[0] == "car"