Genre Linking Blink

🚀 GENRE

GENRE（生成式實體檢索，Generative ENtity REtrieval）系統實現於PyTorch，其相關內容在論文自迴歸實體檢索中有所介紹。簡單來說，GENRE採用序列到序列的方法進行實體檢索（例如實體鏈接），它基於經過微調的BART架構。GENRE通過有約束的束搜索，在輸入文本的條件下生成唯一的實體名稱，以確保只生成有效的標識符。該模型最初在facebookresearch/GENRE倉庫中使用fairseq發佈（transformers模型是通過類似於此的轉換腳本獲得的）。此模型在BLINK的完整訓練集上進行訓練（即基於維基百科的900萬個用於實體消歧的數據點）。

✨ 主要特性

• 採用序列到序列的方法進行實體檢索。
• 基於微調的BART架構。
• 使用有約束的束搜索生成有效標識符。

🚀 快速開始

安裝依賴

你可以使用以下命令安裝必要的庫：

pip install transformers

運行示例代碼

以下是一個用於維基百科頁面消歧的生成示例：

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 = ["Einstein was a [START_ENT] German [END_ENT] physicist."]

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個預測結果（使用有約束的束搜索）：

['Germans',
 'Germany',
 'German Empire',
 'Weimar Republic',
 'Greeks']

📚 詳細文檔

模型信息

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}
}