GENRE

GENRE（生成型实体检索）系统是在 Autoregressive Entity Retrieval 中提出的，并在pytorch中实现。

GENRE使用序列到序列的方法进行实体检索（例如链接），基于经过微调的 BART 架构。GENRE通过生成对输入文本有条件的唯一实体名称来进行检索，使用了受限制的束搜索只生成有效的标识符。该模型首次在 facebookresearch/GENRE 存储库中发布，使用fairseq（使用类似于 this 的转换脚本获得transformers模型）。

该模型是在 KILT 的完整训练集上训练的（即用于事实检查、实体链接、插槽填充、对话、开放域抽取和概括问答的11个数据集）。

BibTeX记录和引用信息

如果您使用来自本存储库的代码，请考虑引用我们的作品。

@inproceedings{decao2020autoregressive,
  title={Autoregressive Entity Retrieval},
  author={Nicola {De Cao} and Gautier Izacard and Sebastian Riedel and Fabio Petroni},
  booktitle={International Conference on Learning Representations},
  url={https://openreview.net/forum?id=5k8F6UU39V},
  year={2021}
}

使用方法

这是用于开放域事实检查的维基百科页面检索的生成示例：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# OPTIONAL: load the prefix tree (trie), you need to additionally download
# https://huggingface.co/facebook/genre-kilt/blob/main/trie.py and 
# https://huggingface.co/facebook/genre-kilt/blob/main/kilt_titles_trie_dict.pkl
# import pickle
# from trie import Trie
# with open("kilt_titles_trie_dict.pkl", "rb") as f:
#     trie = Trie.load_from_dict(pickle.load(f))

tokenizer = AutoTokenizer.from_pretrained("facebook/genre-kilt")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/genre-kilt").eval()

sentences = ["Einstein was a German physicist."]

outputs = model.generate(
    **tokenizer(sentences, return_tensors="pt"),
    num_beams=5,
    num_return_sequences=5,
    # OPTIONAL: use constrained beam search
    # prefix_allowed_tokens_fn=lambda batch_id, sent: trie.get(sent.tolist()),
)

tokenizer.batch_decode(outputs, skip_special_tokens=True)

输出以下前5个预测结果（使用受限束搜索）

['Albert Einstein',
 'Erwin Schrödinger',
 'Werner Bruschke',
 'Werner von Habsburg',
 'Werner von Moltke']