这是 Patrick Lewis、Ethan Perez、Aleksandara Piktus 等人在论文《Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks》(https://arxiv.org/pdf/2005.11401.pdf)中提出的 RAG-Sequence 模型。
该模型是一个uncased模型,这意味着大写字母会被直接转换为小写字母。
该模型由question_encoder(问题编码器)、retriever(检索器)和generator(生成器)组成。检索器从上方链接的wiki_dpr train 数据集中提取相关段落。问题编码器和检索器分别基于 facebook/dpr-question_encoder-single-nq-base 和 facebook/bart-large,并在wiki_dpr问答数据集上以端到端的方式进行了联合微调。
注意:在下面的用法示例中,仅使用了wiki_dpr的dummy(虚拟)检索器,因为完整的legacy(旧版)索引需要超过 75 GB 的内存。
该模型可以按如下方式生成任何事实型问题的答案:
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", retriever=retriever)
input_dict = tokenizer.prepare_seq2seq_batch("how many countries are in europe", return_tensors="pt")
generated = model.generate(input_ids=input_dict["input_ids"])
print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0])
# should give 54 => google says either 44 or 51