语言:

• 英文

标签:

• 检索
• 实体检索
• 命名实体消歧
• 实体消歧
• 命名实体链接
• 实体链接
• 文本生成

GENRE系统

GENRE（生成式实体检索系统）基于论文《自回归实体检索》实现，采用PyTorch框架。简而言之，该系统通过序列到序列的方法进行实体检索（如链接任务），其核心是微调后的BART架构。GENRE通过受约束的束搜索技术生成唯一实体名称，确保输出均为有效标识符。该模型最初由facebookresearch/GENRE仓库使用fairseq发布（当前transformers版本模型通过转换脚本获得，类似于此链接示例）。

本模型基于BLINK完整训练集（即基于维基百科的900万条实体消歧数据点）训练。

引用信息

若使用本仓库代码，请引用我们的工作。

@inproceedings{decao2020autoregressive,
  title={自回归实体检索},
  author={尼古拉·德考 和 高蒂尔·伊扎卡 和 塞巴斯蒂安·里德尔 和 法比奥·佩特罗尼},
  booktitle={国际学习表征会议},
  url={https://openreview.net/forum?id=5k8F6UU39V},
  year={2021}
}

使用示例

以下是维基百科页面消歧的生成示例：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 可选：加载前缀树(trie)，需额外下载
# https://huggingface.co/facebook/genre-linking-blink/blob/main/trie.py 与
# https://huggingface.co/facebook/genre-linking-blink/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-linking-blink")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/genre-linking-blink").eval()

sentences = ["爱因斯坦是[START_ENT]德国[END_ENT]物理学家。"]

outputs = model.generate(
    **tokenizer(sentences, return_tensors="pt"),
    num_beams=5,
    num_return_sequences=5,
    # 可选：使用受约束束搜索
    # prefix_allowed_tokens_fn=lambda batch_id, sent: trie.get(sent.tolist()),
)

tokenizer.batch_decode(outputs, skip_special_tokens=True)

输出前5个预测结果（使用受约束束搜索）：

['德国人',
 '德国',
 '德意志帝国',
 '魏玛共和国',
 '希腊人']