Transformers 文档

RAG

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

RAG

PyTorch TensorFlow FlashAttention

概述

检索增强生成(“RAG”)模型结合了预训练密集检索(DPR)和序列到序列模型的强大功能。RAG 模型检索文档,将它们传递给 seq2seq 模型,然后进行边缘化以生成输出。检索器和 seq2seq 模块从预训练模型初始化,并联合微调,允许检索和生成都适应下游任务。

它基于 Patrick Lewis、Ethan Perez、Aleksandara Piktus、Fabio Petroni、Vladimir Karpukhin、Naman Goyal、Heinrich Küttler、Mike Lewis、Wen-tau Yih、Tim Rocktäschel、Sebastian Riedel、Douwe Kiela 的论文 Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

该论文的摘要如下:

大型预训练语言模型已被证明可以将事实知识存储在其参数中,并在下游 NLP 任务上进行微调时取得最先进的结果。然而,它们访问和精确操纵知识的能力仍然有限,因此在知识密集型任务中,它们的性能落后于特定于任务的架构。此外,为其决策提供出处并更新其世界知识仍然是开放的研究问题。具有可微分访问显式非参数记忆机制的预训练模型可以克服这个问题,但到目前为止,仅针对抽取式下游任务进行了研究。我们探索了一种用于检索增强生成(RAG)的通用微调方法 - 该模型结合了用于语言生成的预训练参数记忆和非参数记忆。我们介绍了 RAG 模型,其中参数记忆是预训练的 seq2seq 模型,非参数记忆是维基百科的密集向量索引,通过预训练的神经检索器访问。我们比较了两种 RAG 公式,一种在整个生成的序列中以相同的检索段落为条件,另一种可以使用每个 token 的不同段落。我们在广泛的知识密集型 NLP 任务上微调和评估了我们的模型,并在三个开放域 QA 任务上设置了最先进的水平,优于参数 seq2seq 模型和特定于任务的检索和抽取架构。对于语言生成任务,我们发现 RAG 模型比最先进的纯参数 seq2seq 基线生成更具体、多样化和基于事实的语言。

此模型由 ola13 贡献。

使用技巧

检索增强生成(“RAG”)模型结合了预训练密集检索(DPR)和 Seq2Seq 模型的强大功能。RAG 模型检索文档 (docs),将它们传递给 seq2seq 模型,然后进行边缘化以生成输出。检索器和 seq2seq 模块从预训练模型初始化,并联合微调,允许检索和生成都适应下游任务。

RagConfig

class transformers.RagConfig

< >

( vocab_size = None is_encoder_decoder = True prefix = None bos_token_id = None pad_token_id = None eos_token_id = None decoder_start_token_id = None title_sep = ' / ' doc_sep = ' // ' n_docs = 5 max_combined_length = 300 retrieval_vector_size = 768 retrieval_batch_size = 8 dataset = 'wiki_dpr' dataset_split = 'train' index_name = 'compressed' index_path = None passages_path = None use_dummy_dataset = False reduce_loss = False label_smoothing = 0.0 do_deduplication = True exclude_bos_score = False do_marginalize = False output_retrieved = False use_cache = True forced_eos_token_id = None dataset_revision = None **kwargs )

参数

  • title_sep (str, 可选, 默认为 " / ") — 在调用 RagRetriever 时,在检索文档的标题和文本之间插入的分隔符。
  • doc_sep (str, 可选, 默认为 " // ") — 在调用 RagRetriever 时,在检索文档的文本和原始输入之间插入的分隔符。
  • n_docs (int, 可选, 默认为 5) — 要检索的文档数量。
  • max_combined_length (int, 可选, 默认为 300) — __call__() 返回的上下文输入的最大长度。
  • retrieval_vector_size (int, 可选, 默认为 768) — 由 RagRetriever 索引的文档嵌入的维度。
  • retrieval_batch_size (int, 可选, 默认为 8) — 检索批次大小,定义为并发发给封装在 RagRetriever 中的 faiss 索引的查询数。
  • dataset (str, 可选, 默认为 "wiki_dpr") — HuggingFace Datasets 中索引数据集的数据集标识符(使用 datasets.list_datasets() 列出所有可用的数据集和 id)。
  • dataset_split (str, 可选, 默认为 "train") — 要加载的 dataset 的哪个拆分。
  • index_name (str, 可选, 默认为 "compressed") — 与 dataset 关联的索引的索引名称。 可以选择 "legacy""exact""compressed" 之一。
  • index_path (str, 可选) — 磁盘上序列化的 faiss 索引的路径。
  • passages_path (str, 可选) — 与 faiss 索引兼容的文本段落的路径。 如果使用 LegacyIndex,则为必需
  • use_dummy_dataset (bool, 可选, 默认为 False) — 是否加载由 dataset 指定的数据集的“虚拟”变体。
  • label_smoothing (float, 可选, 默认为 0.0) — 仅当 return_loss 设置为 True 时相关。 控制损失计算中标签平滑的 epsilon 参数值。 如果设置为 0,则不执行标签平滑。
  • do_marginalize (bool, 可选, 默认为 False) — 如果为 True,则通过使用 torch.nn.functional.log_softmax 对所有文档的 logits 进行边缘化。
  • reduce_loss (bool, 可选, 默认为 False) — 是否使用 torch.Tensor.sum 操作来减少 NLL 损失。
  • do_deduplication (bool, 可选, 默认为 True) — 是否对给定输入的来自不同上下文文档的生成结果进行去重。 如果在分布式后端进行训练时使用,则必须设置为 False
  • exclude_bos_score (bool, 可选, 默认为 False) — 在计算损失时是否忽略 BOS token。
  • output_retrieved(bool, 可选, 默认为 False) — 如果设置为 True,则返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。 有关更多详细信息,请参阅返回的张量。
  • use_cache (bool, 可选, 默认为 True) — 模型是否应返回上次键/值注意力(并非所有模型都使用)。
  • forced_eos_token_id (int, 可选) — 当达到 max_length 时,强制作为最后一个生成的 token 的 token id。 通常设置为 eos_token_id

RagConfig 存储 RagModel 的配置。 配置对象继承自 PretrainedConfig,可用于控制模型输出。 有关更多信息,请阅读 PretrainedConfig 中的文档。

from_question_encoder_generator_configs

< >

( question_encoder_config: PretrainedConfig generator_config: PretrainedConfig **kwargs ) EncoderDecoderConfig

返回:

EncoderDecoderConfig

配置对象的实例

从预训练的编码器模型配置和解码器模型配置实例化 EncoderDecoderConfig(或派生类)。

RagTokenizer

class transformers.RagTokenizer

< >

( question_encoder generator )

Rag 特定输出

transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput

< >

( loss: typing.Optional[torch.FloatTensor] = None logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None retrieved_doc_embeds: typing.Optional[torch.FloatTensor] = None retrieved_doc_ids: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None question_encoder_last_hidden_state: typing.Optional[torch.FloatTensor] = None question_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None )

参数

  • loss (torch.FloatTensor,形状为 (1,)可选,当提供 labels 时返回) — 语言建模损失。
  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。 对于每个词汇表标记,分数可能在所有文档上进行边缘化。
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的分数。
  • past_key_values (List[torch.FloatTensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstorch.FloatTensor 列表,其中每个张量的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算隐藏状态(注意力模块中的键和值),可以用于(参见 past_key_values 输入)加速顺序解码。

  • retrieved_doc_embeds (torch.FloatTensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与 question_encoder_last_hidden_state 一起使用以计算 doc_scores
  • retrieved_doc_ids (torch.LongTensor,形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。
  • context_input_ids (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的输入 ids。
  • context_attention_mask (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的注意力掩码。
  • question_encoder_last_hidden_state (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型的问题编码器池化输出的最后一层输出端的隐藏状态序列。
  • question_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 元组 (每一层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_enc_last_hidden_state (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器的最后一层输出端的隐藏状态序列。
  • generator_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 元组 (每一层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_dec_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 元组 (每一层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_cross_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — torch.FloatTensor 元组 (每一层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均值。

检索增强边缘化模型输出的基类。

transformers.models.rag.modeling_rag.RetrievAugLMOutput

< >

( logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None retrieved_doc_embeds: typing.Optional[torch.FloatTensor] = None retrieved_doc_ids: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None question_encoder_last_hidden_state: typing.Optional[torch.FloatTensor] = None question_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None )

参数

  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。 对于每个词汇表标记,分数可能在所有文档上进行边缘化。
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的分数。
  • past_key_values (List[torch.FloatTensor], optional, returned when use_cache=True is passed or when config.use_cache=True) — 长度为 config.n_layerstorch.FloatTensor 列表,每个张量的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算隐藏状态(注意力模块中的键和值),可以用于(参见 past_key_values 输入)加速顺序解码。

  • retrieved_doc_embeds (形状为 (batch_size, config.n_docs, hidden_size)torch.FloatTensor, 可选, 当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。与 question_encoder_last_hidden_state 一起使用以计算 doc_scores
  • retrieved_doc_ids (形状为 (batch_size, config.n_docs)torch.LongTensor, 可选, 当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。
  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ids。
  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。
  • question_encoder_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor, 可选) — 模型问题编码器池化输出的最后一层的输出端的隐藏状态序列。
  • question_enc_hidden_states (tuple(torch.FloatTensor), 可选, 当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — 形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor 元组(一个用于嵌入的输出,一个用于每一层的输出)。

    问题编码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(torch.FloatTensor), 可选, 当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — 形状为 (batch_size, num_heads, sequence_length, sequence_length)torch.FloatTensor 元组(每一层一个)。

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_enc_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor, 可选) — 模型生成器编码器的最后一层的输出端的隐藏状态序列。
  • generator_enc_hidden_states (tuple(torch.FloatTensor), 可选, 当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — 形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor 元组(一个用于嵌入的输出,一个用于每一层的输出)。

    生成器编码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(torch.FloatTensor), 可选, 当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — 形状为 (batch_size, num_heads, sequence_length, sequence_length)torch.FloatTensor 元组(每一层一个)。

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_dec_hidden_states (tuple(torch.FloatTensor), 可选, 当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — 形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor 元组(一个用于嵌入的输出,一个用于每一层的输出)。

    生成器解码器在每一层输出端的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(torch.FloatTensor), 可选, 当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — 形状为 (batch_size, num_heads, sequence_length, sequence_length)torch.FloatTensor 元组(每一层一个)。

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。

  • generator_cross_attentions (tuple(torch.FloatTensor), 可选, 当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — 形状为 (batch_size, num_heads, sequence_length, sequence_length)torch.FloatTensor 元组(每一层一个)。

    生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均值。

RagRetriever

class transformers.RagRetriever

< >

( config question_encoder_tokenizer generator_tokenizer index = None init_retrieval = True )

参数

  • config (RagConfig) — 与此检索器一起使用的 RAG 模型的配置。包含指示要构建哪个 Index 的参数。您可以使用 config.index_name="custom" 加载您自己的自定义数据集,或者使用来自 datasets 库的规范数据集(默认),例如 config.index_name="wiki_dpr"
  • question_encoder_tokenizer (PreTrainedTokenizer) — 用于对问题进行分词的分词器。它用于解码问题,然后使用 generator_tokenizer。
  • generator_tokenizer (PreTrainedTokenizer) — 用于 RagModel 的生成器部分的分词器。
  • index (Index, optional, defaults to the one defined by the configuration) — 如果指定,则使用此索引代替使用配置构建的索引

用于从向量查询中获取文档的检索器。它检索文档嵌入以及文档内容,并对其进行格式化以与 RagModel 一起使用。

示例

>>> # To load the default "wiki_dpr" dataset with 21M passages from wikipedia (index name is 'compressed' or 'exact')
>>> from transformers import RagRetriever

>>> retriever = RagRetriever.from_pretrained(
...     "facebook/dpr-ctx_encoder-single-nq-base", dataset="wiki_dpr", index_name="compressed"
... )

>>> # To load your own indexed dataset built with the datasets library. More info on how to build the indexed dataset in examples/rag/use_own_knowledge_dataset.py
>>> from transformers import RagRetriever

>>> dataset = (
...     ...
... )  # dataset must be a datasets.Datasets object with columns "title", "text" and "embeddings", and it must have a faiss index
>>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", indexed_dataset=dataset)

>>> # To load your own indexed dataset built with the datasets library that was saved on disk. More info in examples/rag/use_own_knowledge_dataset.py
>>> from transformers import RagRetriever

>>> dataset_path = "path/to/my/dataset"  # dataset saved via *dataset.save_to_disk(...)*
>>> index_path = "path/to/my/index.faiss"  # faiss index saved via *dataset.get_index("embeddings").save(...)*
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/dpr-ctx_encoder-single-nq-base",
...     index_name="custom",
...     passages_path=dataset_path,
...     index_path=index_path,
... )

>>> # To load the legacy index built originally for Rag's paper
>>> from transformers import RagRetriever

>>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", index_name="legacy")

init_retrieval

< >

( )

检索器初始化函数。它将索引加载到内存中。

postprocess_docs

< >

( docs input_strings prefix n_docs return_tensors = None ) tuple(tensors)

参数

  • docs (dict) — 检索到的文档。
  • input_strings (str) — 由 preprocess_query 解码的输入字符串。
  • prefix (str) — 添加到每个输入开头的 前缀,通常用于基于 T5 的模型。

返回:

tuple(tensors)

一个由两个元素组成的元组:上下文相关的 input_ids 和兼容的 attention_mask

后处理检索到的 docs 并将它们与 input_strings 结合起来。

retrieve

< >

( question_hidden_states: ndarray n_docs: int ) Tuple[np.ndarray, np.ndarray, List[dict]]

参数

  • question_hidden_states (np.ndarray,形状为 (batch_size, vector_size)) — 用于检索的查询向量批次。
  • n_docs (int) — 每个查询检索的文档数量。

返回:

Tuple[np.ndarray, np.ndarray, List[dict]]

包含以下对象的元组

  • retrieved_doc_embeds (np.ndarray,形状为 (batch_size, n_docs, dim)) — 每个查询检索到的文档的检索嵌入。
  • doc_ids (np.ndarray,形状为 (batch_size, n_docs)) — 索引中文档的 ID
  • doc_dicts (List[dict]): 每个查询的 retrieved_doc_embeds 示例。

检索指定 question_hidden_states 的文档。

Pytorch
隐藏 Pytorch 内容

RagModel

class transformers.RagModel

< >

( config: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )

参数

  • config (RagConfig) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重,只会加载配置。查看 from_pretrained() 方法加载模型权重。
  • question_encoder (PreTrainedModel) — 与 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 检索器类,封装了 faiss 索引,用于查询以获取当前输入的上下文文档。

RagModel forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration

该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder 和任何带有语言模型头的 seq2seq 模型作为 generator。它已经使用 DPRQuestionEncoder 作为 question_encoderBartForConditionalGenerationT5ForConditionalGeneration 作为 generator 进行了测试。

此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.FloatTensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.FloatTensor]]] = None doc_scores: typing.Optional[torch.FloatTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None n_docs: typing.Optional[int] = None ) transformers.models.rag.modeling_rag.RetrievAugLMOutputtuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor,形状为 (batch_size, sequence_length)) — 词汇表中输入序列标记的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。

    什么是输入 ID?

  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — 避免在 padding 标记索引上执行 attention 的 Mask。在 [0, 1] 中选择的 Mask 值:

    • 1 表示未被 Mask 的标记,
    • 0 表示被 Mask 的标记。

    什么是 attention mask?

  • encoder_outputs (tuple(tuple(torch.FloatTensor), 可选) — 元组由 (generator_enc_last_hidden_state, 可选: generator_enc_hidden_states, 可选: generator_enc_attentions) 组成。形状为 (batch_size, n_docs * sequence_length, hidden_size)generator_enc_last_hidden_state 是生成器编码器最后一层输出端的 hidden-states 序列。

    由 (RagModel) 模型在解码期间使用。

  • decoder_input_ids (torch.LongTensor,形状为 (batch_size, target_sequence_length)可选) — 为生成任务提供。默认为 None,根据您在 RAG 实例中使用的生成器模型的说明构建。
  • decoder_attention_mask (torch.BoolTensor,形状为 (batch_size, target_sequence_length)可选) — 默认行为:生成一个忽略 decoder_input_ids 中 pad 标记的 tensor。默认情况下,也将使用因果 mask。
  • past_key_values (tuple(tuple(torch.FloatTensor))) — 元组由两个元素组成:RAG 模型的 encoder_outputs(参见 encoder_outputs)和底层生成器的 past_key_values。可用于加速解码。在 (RagTokenForGeneration) 模型解码期间使用 past_key_values
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。如果模型没有使用 retriever 初始化,则必须将 doc_scores 提供给 forward 传递。 doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,有关更多信息,请参见示例。
  • context_input_ids (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 进行后处理的输入 ID。如果模型没有使用 retriever 初始化,则必须将 context_input_ids 提供给 forward 传递。 context_input_ids__call__() 返回。
  • context_attention_mask (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 进行后处理的 Attention mask。如果模型没有使用 retriever 初始化,则必须将 context_attention_mask 提供给 forward 传递。 context_attention_mask__call__() 返回。
  • use_cache (bool可选,默认为 True) — 如果设置为 True,则返回 past_key_values 键值状态,并可用于加速解码(参见 past_key_values)。
  • output_attentions (bool可选) — 是否返回所有 attention 层的 attentions tensor。有关更多详细信息,请参见返回的 tensor 下的 attentions
  • output_hidden_states (bool可选) — 是否返回所有层的 hidden states。有关更多详细信息,请参见返回的 tensor 下的 hidden_states
  • output_retrieved(bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。有关更多详细信息,请参见返回的 tensor。
  • n_docs (int可选,默认为 `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。

返回:

transformers.models.rag.modeling_rag.RetrievAugLMOutputtuple(torch.FloatTensor)

一个 transformers.models.rag.modeling_rag.RetrievAugLMOutput 或一个 torch.FloatTensor 元组(如果传递了 return_dict=False 或当 config.return_dict=False 时),包括各种元素,具体取决于配置 (RagConfig) 和输入。

  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。分数可能是针对每个词汇表标记对所有文档进行边缘化的结果。

  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • past_key_values (List[torch.FloatTensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstorch.FloatTensor 列表,其中每个 tensor 的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • retrieved_doc_embeds (torch.FloatTensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (torch.LongTensor,形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。

  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 id。

  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

  • question_encoder_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。

  • question_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 模型生成器编码器的最后一层输出的隐藏状态序列。

  • generator_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_cross_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均。

RagModel forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> from transformers import AutoTokenizer, RagRetriever, RagModel
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = RagModel.from_pretrained("facebook/rag-token-base", retriever=retriever)

>>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
>>> outputs = model(input_ids=inputs["input_ids"])

RagSequenceForGeneration

class transformers.RagSequenceForGeneration

< >

( config: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )

参数

  • config (RagConfig) — 带有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重,仅加载配置。查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (PreTrainedModel) — 与由 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。

RagSequenceForGeneration 的 forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

RAG 序列模型实现。它在前向传播中执行 RAG 序列特定的边缘化。

RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration

该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder 和任何带有语言模型头的 seq2seq 模型作为 generator。它已经使用 DPRQuestionEncoder 作为 question_encoderBartForConditionalGenerationT5ForConditionalGeneration 作为 generator 进行了测试。

此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None exclude_bos_score: typing.Optional[bool] = None reduce_loss: typing.Optional[bool] = None labels: typing.Optional[torch.LongTensor] = None n_docs: typing.Optional[int] = None **kwargs ) transformers.models.rag.modeling_rag.RetrievAugLMMarginOutputtuple(torch.FloatTensor)

参数

  • input_ids (形状为 (batch_size, sequence_length)torch.LongTensor) — 输入序列 token 在词汇表中的索引。RagConfig,用于初始化模型,指定要使用的生成器,还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。

    什么是输入 IDs?

  • attention_mask (形状为 (batch_size, sequence_length)torch.Tensor可选) — 用于避免在 padding token 索引上执行注意力的掩码。在 [0, 1] 中选择的掩码值:

    • 1 表示 未被掩盖 的 token,
    • 0 表示 被掩盖 的 token。

    什么是注意力掩码?

  • encoder_outputs (tuple(tuple(torch.FloatTensor)可选) — 元组由 (generator_enc_last_hidden_state可选generator_enc_hidden_states可选generator_enc_attentions) 组成。形状为 (batch_size, n_docs * sequence_length, hidden_size)generator_enc_last_hidden_state 是生成器编码器最后一层输出的隐藏状态序列。

    由 (RagModel) 模型在解码期间使用。

  • decoder_input_ids (形状为 (batch_size, target_sequence_length)torch.LongTensor可选) — 为生成任务提供。默认为 None,根据您与 RAG 实例一起使用的生成器模型的说明构建。
  • decoder_attention_mask (形状为 (batch_size, target_sequence_length)torch.BoolTensor可选) — 默认行为:生成一个忽略 decoder_input_ids 中的 pad token 的 tensor。默认情况下也将使用因果掩码。
  • past_key_values (tuple(tuple(torch.FloatTensor))) — 元组由两个元素组成:RAG 模型的 encoder_outputs (参见 encoder_outputs) 和底层生成器的 past_key_values。可用于加速解码。 past_key_values 在 (RagTokenForGeneration) 模型解码期间使用。
  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入 (参见 retrieved_doc_embeds) 和 question_encoder_last_hidden_state 之间的分数。如果模型未用 retriever 初始化,则必须将 doc_scores 提供给前向传播。 doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,有关更多信息,请参见示例。
  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。如果模型未用 retriever 初始化,则必须将 context_input_ids 提供给前向传播。 context_input_ids__call__() 返回。
  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。如果模型未用 retriever 初始化,则必须将 context_attention_mask 提供给前向传播。 context_attention_mask__call__() 返回。
  • use_cache (bool可选,默认为 True) — 如果设置为 True,则返回 past_key_values 键值状态,并可用于加速解码 (参见 past_key_values)。
  • output_attentions (bool可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息,请参见返回的张量下的 attentions
  • output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。 更多细节请参阅返回张量下的 hidden_states
  • output_retrieved(bool, 可选) — 是否返回 retrieved_doc_embeds, retrieved_doc_ids, context_input_idscontext_attention_mask。 更多细节请参阅返回张量。
  • n_docs (int, 可选, 默认为 `config.n_docs`) — 检索的文档数量和/或生成答案的文档数量。
  • exclude_bos_score (bool, 可选) — 仅当传递 labels 时相关。 如果为 True,则在计算损失时忽略 BOS 标记的分数。
  • reduce_loss (bool, 可选) — 仅当传递 labels 时相关。 如果为 True,则使用 torch.Tensor.sum 操作来减少 NLL 损失。
  • kwargs (Dict[str, any], 可选, 默认为 {}) — 旧版字典,模型需要它才能使用 generate() 函数。

返回:

transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput or tuple(torch.FloatTensor)

一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组 (如果传递了 return_dict=False 或者当 config.return_dict=False 时),包含各种元素,具体取决于配置 (RagConfig) 和输入。

  • loss (形状为 (1,)torch.FloatTensor, 可选, 当提供 labels 时返回) — 语言建模损失。

  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。分数可能是针对每个词汇表标记对所有文档进行边缘化的结果。

  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • past_key_values (List[torch.FloatTensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstorch.FloatTensor 列表,其中每个 tensor 的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • retrieved_doc_embeds (torch.FloatTensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (torch.LongTensor,形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。

  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 id。

  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

  • question_encoder_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。

  • question_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 模型生成器编码器的最后一层输出的隐藏状态序列。

  • generator_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_cross_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均。

RagSequenceForGeneration 的 forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> from transformers import AutoTokenizer, RagRetriever, RagSequenceForGeneration
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)

>>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
>>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt")
>>> input_ids = inputs["input_ids"]
>>> labels = targets["input_ids"]
>>> outputs = model(input_ids=input_ids, labels=labels)

>>> # or use retriever separately
>>> model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", use_dummy_dataset=True)
>>> # 1. Encode
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
>>> doc_scores = torch.bmm(
...     question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)
... ).squeeze(1)
>>> # 3. Forward to generator
>>> outputs = model(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
...     decoder_input_ids=labels,
... )

generate

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None do_deduplication: typing.Optional[bool] = None num_return_sequences: typing.Optional[int] = None num_beams: typing.Optional[int] = None n_docs: typing.Optional[int] = None **model_kwargs ) torch.LongTensor of shape (batch_size * num_return_sequences, sequence_length)

参数

  • input_ids (形状为 (batch_size, sequence_length)torch.LongTensor, 可选) — 用作生成的提示的序列。 如果未传递 input_ids,则必须提供 context_input_ids
  • attention_mask (形状为 (batch_size, sequence_length)torch.Tensor, 可选) — 用于避免对填充标记索引执行注意力的掩码。 在 [0, 1] 中选择的掩码值:

    • 1 表示未被掩盖的标记,
    • 0 表示被掩盖的标记。

    什么是注意力掩码?

  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 预处理后的输入 ID。
  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 预处理后的注意力掩码。

    如果模型未用 retriever 初始化,或者未给出 input_ids,则必须将 context_input_idscontext_attention_mask 提供给前向传递。 它们由 __call__() 返回。

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入 (参见 retrieved_doc_embeds) 与 question_encoder_last_hidden_state 之间的分数。

    如果模型未用 retriever 初始化,或者未给出 input_ids,则必须将 doc_scores 提供给前向传递。 doc_scores__call__() 返回。

  • do_deduplication (bool, 可选) — 是否为给定输入去重来自不同上下文文档的生成结果。 如果在使用分布式后端进行训练时使用,则必须设置为 False
  • num_return_sequences(int, 可选, 默认为 1) — 批次中每个元素独立计算的返回序列数。 请注意,这不是我们传递给 generator[generate()](/docs/transformers/v4.50.0/en/main_classes/text_generation#transformers.GenerationMixin.generate) 函数的值,我们在该函数中将 num_return_sequences 设置为 num_beams
  • num_beams (int, 可选, 默认为 1) — 束搜索的束数量。 1 表示不进行束搜索。
  • n_docs (int, 可选, 默认为 config.n_docs) — 检索的文档数量和/或生成答案的文档数量。
  • kwargs (Dict[str, Any], 可选) — 其他 kwargs 将传递给 generate()

返回:

torch.LongTensor of shape (batch_size * num_return_sequences, sequence_length)

生成的序列。 第二个维度(序列长度)要么等于 max_length,要么更短,如果所有批次由于 eos_token_id 而提前完成。

实现 RAG 序列“彻底”解码。 阅读 generate()` 文档以获取有关如何设置其他生成输入参数的更多信息。

RagTokenForGeneration

transformers.RagTokenForGeneration

< >

( config: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )

参数

  • config (RagConfig) — 带有模型所有参数的模型配置类。 使用配置文件初始化不会加载与模型关联的权重,仅加载配置。 查看 from_pretrained() 方法来加载模型权重。
  • question_encoder (PreTrainedModel) — 与 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,该索引被查询以获取当前输入的上下文文档。

RagTokenForGeneration 的前向方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

RAG-token 模型实现。 它在前向传递中执行 RAG-token 特定的边缘化。

RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration

该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder 和任何带有语言模型头的 seq2seq 模型作为 generator。它已经使用 DPRQuestionEncoder 作为 question_encoderBartForConditionalGenerationT5ForConditionalGeneration 作为 generator 进行了测试。

此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.FloatTensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None do_marginalize: typing.Optional[bool] = None reduce_loss: typing.Optional[bool] = None labels: typing.Optional[torch.LongTensor] = None n_docs: typing.Optional[int] = None **kwargs ) transformers.models.rag.modeling_rag.RetrievAugLMMarginOutputtuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor,形状为 (batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定了要使用的生成器,并且还指定了兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。

    什么是输入 IDs?

  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — Mask,用于避免对 padding token 索引执行 attention。Mask 值在 [0, 1] 中选择:

    • 1 表示 **未被 mask** 的 tokens,
    • 0 表示 **已被 mask** 的 tokens。

    什么是 attention masks?

  • encoder_outputs (tuple(tuple(torch.FloatTensor), 可选) — Tuple 由 (generator_enc_last_hidden_state, 可选: generator_enc_hidden_states, 可选: generator_enc_attentions) 组成。形状为 (batch_size, n_docs * sequence_length, hidden_size)generator_enc_last_hidden_state 是生成器 encoder 最后一层输出的 hidden-states 序列。

    由 (RagModel) 模型在解码期间使用。

  • decoder_input_ids (torch.LongTensor,形状为 (batch_size, target_sequence_length)可选) — 为生成任务提供。 默认值为 None,根据您在 RAG 实例中使用的生成器模型的说明构建。
  • decoder_attention_mask (torch.BoolTensor,形状为 (batch_size, target_sequence_length)可选) — 默认行为:生成一个 tensor,该 tensor 忽略 decoder_input_ids 中的 pad tokens。因果 mask 也将默认使用。
  • past_key_values (tuple(tuple(torch.FloatTensor))) — Tuple 由两个元素组成:RAG 模型的 encoder_outputs (参见 encoder_outputs) 和底层生成器的 past_key_values。可用于加速解码。past_key_values 在 (RagTokenForGeneration) 模型解码期间使用。
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入 (参见 retrieved_doc_embeds) 和 question_encoder_last_hidden_state 之间的分数。如果模型没有使用 retriever 初始化,则必须将 doc_scores 提供给 forward 传递。doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,有关更多信息,请参见示例。
  • context_input_ids (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的输入 IDs。如果模型没有使用 retriever 初始化,则必须将 `context_input_ids 提供给 forward 传递。 context_input_ids__call__() 返回。
  • context_attention_mask (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的 Attention mask。如果模型没有使用 retriever 初始化,则必须将 context_attention_mask 提供给 forward 传递。 context_attention_mask__call__() 返回。
  • use_cache (bool可选,默认为 True) — 如果设置为 True,则返回 past_key_values 键值状态,并且可以用于加速解码 (参见 past_key_values)。
  • output_attentions (bool可选) — 是否返回所有 attention 层的 attentions tensors。 有关更多详细信息,请参见返回的 tensors 下的 attentions
  • output_hidden_states (bool可选) — 是否返回所有层的 hidden states。 有关更多详细信息,请参见返回的 tensors 下的 hidden_states
  • output_retrieved(bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。有关更多详细信息,请参见返回的 tensors。
  • n_docs (int可选,默认为 `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • do_marginalize (bool可选) — 如果为 True,则通过使用 torch.nn.functional.log_softmax 对所有文档的 logits 进行边缘化。
  • reduce_loss (bool可选) — 仅在传递 labels 时相关。如果为 True,则使用 torch.Tensor.sum 操作来减少 NLL 损失。
  • kwargs (Dict[str, any]可选,默认为 {}) — 遗留字典,模型需要它才能使用 generate() 函数。

返回:

transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput or tuple(torch.FloatTensor)

一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组 (如果传递了 return_dict=False 或者当 config.return_dict=False 时),包含各种元素,具体取决于配置 (RagConfig) 和输入。

  • loss (形状为 (1,)torch.FloatTensor, 可选, 当提供 labels 时返回) — 语言建模损失。

  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。分数可能是针对每个词汇表标记对所有文档进行边缘化的结果。

  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • past_key_values (List[torch.FloatTensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstorch.FloatTensor 列表,其中每个 tensor 的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • retrieved_doc_embeds (torch.FloatTensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (torch.LongTensor,形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。

  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 id。

  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

  • question_encoder_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。

  • question_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor可选) — 模型生成器编码器的最后一层输出的隐藏状态序列。

  • generator_enc_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(torch.FloatTensor)可选,当传递 output_hidden_states=Trueconfig.output_hidden_states=True 时返回) — torch.FloatTensor 的元组 (每个输出层对应一个,包括嵌入层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_cross_attentions (tuple(torch.FloatTensor)可选,当传递 output_attentions=Trueconfig.output_attentions=True 时返回) — torch.FloatTensor 的元组 (每一层对应一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均。

RagTokenForGeneration 的前向方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> from transformers import AutoTokenizer, RagRetriever, RagTokenForGeneration
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)

>>> inputs = tokenizer("How many people live in Paris?", return_tensors="pt")
>>> targets = tokenizer(text_target="In Paris, there are 10 million people.", return_tensors="pt")
>>> input_ids = inputs["input_ids"]
>>> labels = targets["input_ids"]
>>> outputs = model(input_ids=input_ids, labels=labels)

>>> # or use retriever separately
>>> model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
>>> # 1. Encode
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.detach().numpy(), return_tensors="pt")
>>> doc_scores = torch.bmm(
...     question_hidden_states.unsqueeze(1), docs_dict["retrieved_doc_embeds"].float().transpose(1, 2)
... ).squeeze(1)
>>> # 3. Forward to generator
>>> outputs = model(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
...     decoder_input_ids=labels,
... )

>>> # or directly generate
>>> generated = model.generate(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)

generate

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None n_docs: typing.Optional[int] = None generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None prefix_allowed_tokens_fn: typing.Callable[[int, torch.Tensor], typing.List[int]] = None logits_processor: typing.Optional[transformers.generation.logits_process.LogitsProcessorList] = [] stopping_criteria: typing.Optional[transformers.generation.stopping_criteria.StoppingCriteriaList] = [] **kwargs ) torch.LongTensor,形状为 (batch_size * num_return_sequences, sequence_length)

参数

  • input_ids (torch.LongTensor,形状为 (batch_size, sequence_length)可选) — 用作生成的 prompt 的序列。如果未传递 input_ids,则必须提供 context_input_ids
  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — Mask,用于避免对 padding token 索引执行 attention。Mask 值在 [0, 1] 中选择:

    • 1 表示 **未被 mask** 的 tokens,
    • 0 表示 **已被 mask** 的 tokens。

    什么是 attention masks?

  • context_input_ids (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的输入 IDs。

    如果模型没有使用 retriever 初始化,则必须将 context_input_ids 提供给 forward 传递。 context_input_ids__call__() 返回。

  • context_attention_mask (torch.LongTensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理的 Attention mask。

    如果模型没有使用 retriever 初始化,则必须将 context_input_ids 提供给 forward 传递。 context_input_ids__call__() 返回。

  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入 (参见 retrieved_doc_embeds) 和 question_encoder_last_hidden_state 之间的分数。

    如果模型没有使用 retriever 初始化,则必须将 context_input_ids 提供给 forward 传递。 context_input_ids__call__() 返回。

  • n_docs (int可选,默认为 config.n_docs) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • generation_config (~generation.GenerationConfig, 可选) — 用作生成调用的基本参数化的生成配置。传递给 generate 的与 generation_config 属性匹配的 **kwargs 将覆盖它们。如果未提供 generation_config,则将使用默认值,其加载优先级如下:1) 来自 generation_config.json 模型文件(如果存在);2) 来自模型配置。请注意,未指定的参数将继承 GenerationConfig 的默认值,应查阅其文档以参数化生成。
  • prefix_allowed_tokens_fn (Callable[[int, torch.Tensor], List[int]], 可选) — 如果提供,此函数将在每一步将束搜索约束为仅允许的 token。如果未提供,则不应用约束。此函数接受 2 个参数 inputs_ids 和批次 ID batch_id。它必须返回一个列表,其中包含用于下一步生成的允许 token,这些 token 以先前生成的 token inputs_ids 和批次 ID batch_id 为条件。此参数对于以 prefix 为条件的约束生成很有用,如 Autoregressive Entity Retrieval 中所述。
  • logits_processor (LogitsProcessorList, 可选) — 补充从参数和模型配置构建的默认 logits 处理器的自定义 logits 处理器。如果传递的 logit 处理器已使用参数或模型配置创建,则会抛出错误。
  • stopping_criteria (StoppingCriteriaList, 可选) — 补充从参数和模型配置构建的默认停止标准的自定义停止标准。如果传递的停止标准已使用参数或模型配置创建,则会抛出错误。
  • kwargs (Dict[str, Any], 可选) — generate_config 的特别参数化和/或将转发到模型的 forward 函数的其他模型特定 kwargs。

返回:

torch.LongTensor of shape (batch_size * num_return_sequences, sequence_length)

生成的序列。第二个维度(sequence_length)要么等于 max_length,要么如果所有批次由于 eos_token_id 提前完成,则更短。

实现 RAG token 解码。

TensorFlow
隐藏 TensorFlow 内容

TFRagModel

class transformers.TFRagModel

< >

( config: Optional[PretrainedConfig] = None question_encoder: Optional[TFPreTrainedModel] = None generator: Optional[TFPreTrainedModel] = None retriever: Optional[RagRetriever] = None load_weight_prefix: Optional[str] = None **kwargs )

参数

  • config (RagConfig) — 模型配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重,仅加载配置。查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (TFPreTrainedModel) — 与 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (TFPreTrainedModel) — 用作 RAG 架构中生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 一个检索器类,封装了 faiss 索引,该索引被查询以获取当前输入的上下文文档。

TFRagModel 的 forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在正向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,生成器可以是任何seq2seq模型,最好是 TFBartForConditionalGeneration

该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中——有关更多详细信息,请参见示例。该模型兼容任何自编码模型作为 question_encoder,以及任何带有语言模型头的 seq2seq 模型作为 generator。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder,使用 TFBartForConditionalGeneration 作为 generator

此模型继承自 TFPreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 TensorFlow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与一般用法和行为相关的所有事项。

该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。

call

< >

( input_ids: TFModelInputType | None = None attention_mask: np.ndarray | tf.Tensor | None = None encoder_outputs: np.ndarray | tf.Tensor | None = None decoder_input_ids: np.ndarray | tf.Tensor | None = None decoder_attention_mask: np.ndarray | tf.Tensor | None = None past_key_values: Tuple[Tuple[Union[np.ndarray, tf.Tensor]]] | None = None doc_scores: np.ndarray | tf.Tensor | None = None context_input_ids: np.ndarray | tf.Tensor | None = None context_attention_mask: np.ndarray | tf.Tensor | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None output_retrieved: bool | None = None n_docs: int | None = None return_dict: bool | None = None training: bool = False **kwargs transformers.models.rag.modeling_tf_rag.TFRetrievAugLMOutput or tuple(tf.Tensor)

参数

  • input_ids (tf.Tensor of shape (batch_size, sequence_length)) — 词汇表中输入序列 token 的索引。RagConfig,用于初始化模型,指定要使用的生成器,还指定兼容的生成器 tokenizer。
  • attention_mask (tf.Tensor of shape (batch_size, sequence_length), 可选) — 用于避免对 padding token 索引执行 attention 的掩码。掩码值在 [0, 1] 中选择:

    • 1 代表未被掩码的 token,
    • 0 代表被掩码的 token。

    什么是 attention masks?

  • encoder_outputs (tuple(tuple(tf.Tensor), 可选) — 元组由 (generator_enc_last_hidden_state, 可选: generator_enc_hidden_states, 可选: generator_enc_attentions) 组成。generator_enc_last_hidden_state 的形状为 (batch_size, n_docs * sequence_length, hidden_size),是生成器编码器最后一层输出的 hidden-state 序列。

    由 (TFRagModel) 模型在解码期间使用。

  • decoder_input_ids (tf.Tensor of shape (batch_size, target_sequence_length), 可选) — 为生成任务提供。默认为 None,按照您与 RAG 实例一起使用的生成器模型的说明构建。
  • decoder_attention_mask (torch.BoolTensor of shape (batch_size, target_sequence_length), 可选) — 默认行为:生成一个忽略 decoder_input_ids 中 pad token 的 tensor。默认情况下还将使用因果掩码。
  • past_key_values (tuple(tuple(tf.Tensor))) — 元组由两个元素组成:RAG 模型的 encoder_outputs(参见 encoder_outputs)和底层生成器的 past_key_values。可用于加速解码。 past_key_values 在 (RagTokenForGeneration) 模型解码期间使用。
  • doc_scores (tf.Tensor of shape (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。如果模型未用 retriever 初始化,则必须将 doc_scores 提供给 forward 传递。 doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算得出,有关更多信息,请参见示例。
  • context_input_ids (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 从检索到的文档和问题编码器 input_ids 经过检索器后处理的输入 ID。

    如果模型没有用 retriever 初始化,则必须将 context_input_ids 提供给前向传播。 context_input_ids__call__() 返回。 context_attention_mask (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回): 从检索到的文档和问题编码器 input_ids 经过检索器后处理的注意力掩码。

    如果模型没有用 retriever 初始化,则必须将 context_attention_mask 提供给前向传播。 context_attention_mask__call__() 返回。

  • use_cache (bool可选,默认为 True) — 如果设置为 True,则返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • output_attentions (bool可选) — 是否返回所有注意力层的注意力张量。 有关更多详细信息,请参见返回的张量下的 attentions
  • output_hidden_states (bool可选) — 是否返回所有层的隐藏状态。 有关更多详细信息,请参见返回的张量下的 hidden_states
  • output_retrieved(bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。 有关更多详细信息,请参见返回的张量。
  • return_dict (bool可选) — 是否返回 TFRetrievAugLMOutput 而不是普通元组。
  • n_docs (int可选,默认为 config.n_docs) — 要检索的文档数和/或要为其生成答案的文档数。

返回:

transformers.models.rag.modeling_tf_rag.TFRetrievAugLMOutputtuple(tf.Tensor)

一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMOutputtf.Tensor 元组(如果传递 return_dict=False 或当 config.return_dict=False 时),其中包含各种元素,具体取决于配置 (RagConfig) 和输入。

  • logits (tf.Tensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。 该分数可能针对每个词汇标记在所有文档中进行边缘化处理。

  • past_key_values (List[tf.Tensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstf.Tensor 列表,其中每个张量的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • retrieved_doc_embeds (tf.Tensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (tf.Tensor,形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。

  • context_input_ids (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 从检索到的文档和问题编码器 input_ids 经过检索器后处理的输入 id。

  • context_attention_mask (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length)可选,当 output_retrieved=True 时返回) — 从检索到的文档和问题编码器 input_ids 经过检索器后处理的注意力掩码。

  • question_encoder_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。

  • question_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。

  • generator_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

TFRagModel 的 forward 方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> from transformers import AutoTokenizer, RagRetriever, TFRagModel
>>> import torch

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagModel.from_pretrained("facebook/rag-token-base", retriever=retriever, from_pt=True)

>>> input_dict = tokenizer.prepare_seq2seq_batch(
...     "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> input_ids = input_dict["input_ids"]
>>> outputs = model(input_ids)

TFRagSequenceForGeneration

class transformers.TFRagSequenceForGeneration

< >

( config: Optional[PretrainedConfig] = None question_encoder: Optional[TFPreTrainedModel] = None generator: Optional[TFPreTrainedModel] = None retriever: Optional[RagRetriever] = None **kwargs )

参数

  • config (RagConfig) — 模型配置类,其中包含模型的所有参数。 使用配置文件初始化不会加载与模型关联的权重,而只会加载配置。 查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (TFPreTrainedModel) — 与 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (TFPreTrainedModel) — 用作 RAG 架构中生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 检索器类,封装了 faiss 索引,该索引被查询以获取当前输入的上下文文档。

TFRagSequenceForGeneration 前向方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

TF RAG 序列模型实现。 它在前向传播中执行 RAG 序列特定的边缘化处理。

RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在正向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,生成器可以是任何seq2seq模型,最好是 TFBartForConditionalGeneration

该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中——有关更多详细信息,请参见示例。该模型兼容任何自编码模型作为 question_encoder,以及任何带有语言模型头的 seq2seq 模型作为 generator。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder,使用 TFBartForConditionalGeneration 作为 generator

此模型继承自 TFPreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 TensorFlow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与一般用法和行为相关的所有事项。

该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。

call

< >

( input_ids: TFModelInputType | None = None attention_mask: np.ndarray | tf.Tensor | None = None decoder_input_ids: np.ndarray | tf.Tensor | None = None decoder_attention_mask: np.ndarray | tf.Tensor | None = None encoder_outputs: np.ndarray | tf.Tensor | None = None past_key_values: Optional[Tuple[Tuple[Union[np.ndarray, tf.Tensor]]]] = None doc_scores: np.ndarray | tf.Tensor | None = None context_input_ids: np.ndarray | tf.Tensor | None = None context_attention_mask: np.ndarray | tf.Tensor | None = None use_cache: Optional[bool] = None output_attentions: Optional[bool] = None output_hidden_states: Optional[bool] = None output_retrieved: Optional[bool] = None n_docs: Optional[int] = None exclude_bos_score: Optional[bool] = None labels: np.ndarray | tf.Tensor | None = None reduce_loss: Optional[bool] = None return_dict: Optional[bool] = None training: bool = False **kwargs ) transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutputtuple(tf.Tensor)

参数

  • input_ids (tf.Tensor,形状为 (batch_size, sequence_length)) — 词汇表中输入序列标记的索引。 RagConfig,用于初始化模型,指定要使用的生成器,它还指定兼容的生成器 tokenizer。 使用该 tokenizer 类来获取索引。
  • attention_mask (tf.Tensor,形状为 (batch_size, sequence_length)可选) — 掩码,以避免在填充令牌索引上执行注意力机制。 在 [0, 1] 中选择的掩码值:

    • 1 表示未被掩盖的标记,
    • 0 表示被掩盖的标记。

    什么是注意力掩码?

  • encoder_outputs (tuple(tuple(tf.Tensor)可选) — 元组由 (generator_enc_last_hidden_state, 可选: generator_enc_hidden_states, 可选: generator_enc_attentions) 组成。 形状为 (batch_size, n_docs * sequence_length, hidden_size)generator_enc_last_hidden_state 是生成器编码器最后一层输出处的隐藏状态序列。

    由 (TFRagModel) 模型在解码期间使用。

  • decoder_input_ids (tf.Tensor,形状为 (batch_size, target_sequence_length)可选) — 为生成任务提供。 默认为 None,按照您正在与 RAG 实例一起使用的生成器模型的说明构建。
  • decoder_attention_mask (torch.BoolTensor,形状为 (batch_size, target_sequence_length)可选) — 默认行为:生成一个张量,该张量忽略 decoder_input_ids 中的 pad 标记。 默认情况下,也将使用因果掩码。
  • past_key_values (tuple(tuple(tf.Tensor))) — 元组由两个元素组成:RAG 模型的 encoder_outputs(请参阅 encoder_outputs)和底层生成器的 past_key_values。 可用于加速解码。 past_key_values 在 (RagTokenForGeneration) 模型在解码期间使用。
  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。 如果模型没有用 retriever 初始化,则必须将 doc_scores 提供给前向传播。 doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算得出,请参阅示例以获取更多信息。
  • context_input_ids (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。

    如果模型没有用 retriever 初始化,则必须将 context_input_ids 提供给前向传递。 context_input_ids__call__() 返回。 context_attention_mask (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True): 检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

    如果模型没有用 retriever 初始化,则必须将 context_attention_mask 提供给前向传递。 context_attention_mask__call__() 返回。

  • use_cache (bool, optional, defaults to True) — 如果设置为 True,则返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • output_attentions (bool, optional) — 是否返回所有注意力层的注意力张量。 有关更多详细信息,请参阅返回张量下的 attentions
  • output_hidden_states (bool, optional) — 是否返回所有层的隐藏状态。 有关更多详细信息,请参阅返回张量下的 hidden_states
  • output_retrieved(bool, optional) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。 有关更多详细信息,请参阅返回张量。
  • return_dict (bool, optional) — 是否返回 TFRetrievAugLMOutput 而不是普通元组。
  • n_docs (int, optional, defaults to `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • exclude_bos_score (bool, optional) — 仅在传递 labels 时相关。 如果为 True,则在计算损失时忽略 BOS 标记的分数。
  • labels (tf.Tensornp.ndarray of shape (batch_size, sequence_length), optional) — 用于根据 Rag-Sequence 模型公式计算交叉熵分类损失的标签。 有关 Rag-Sequence 公式的详细信息,请参阅 https://arxiv.org/pdf/2005.11401.pdf 第 2.1 节。 索引应在 [0, ..., config.vocab_size - 1] 中。
  • reduce_loss (bool, optional) — 仅在传递 labels 时相关。 如果为 True,则使用 tf.Tensor.sum 操作减少 NLL 损失。
  • kwargs (Dict[str, any], optional, defaults to {}) — 旧版字典,模型需要它才能使用 generate() 函数。

返回:

transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutputtuple(tf.Tensor)

一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutputtf.Tensor 的元组(如果传递了 return_dict=False 或当 config.return_dict=False 时),其中包含各种元素,具体取决于配置 (RagConfig) 和输入。

  • loss (tf.Tensor of shape (1,), optional, returned when labels is provided) — 语言建模损失。

  • logits (tf.Tensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。 该分数可能针对每个词汇标记在所有文档中进行边缘化处理。

  • past_key_values (List[tf.Tensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstf.Tensor 列表,其中每个张量的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • retrieved_doc_embeds (tf.Tensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (tf.Tensor (int32) of shape (batch_size, config.n_docs), optional, returned when output_retrieved=True) — 检索器检索到的嵌入文档的索引。

  • context_input_ids (tf.Tensor(int32) of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。

  • context_attention_mask (tf.Tensor (int32) of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

  • question_encoder_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。

  • question_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。

  • generator_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

TFRagSequenceForGeneration 前向方法,覆盖了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> from transformers import AutoTokenizer, RagRetriever, TFRagSequenceForGeneration

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagSequenceForGeneration.from_pretrained(
...     "facebook/rag-sequence-nq", retriever=retriever, from_pt=True
... )

>>> input_dict = tokenizer.prepare_seq2seq_batch(
...     "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> outputs = model(input_dict, output_retrieved=True)

>>> # or use retriever separately
>>> # 1. Encode
>>> input_ids = input_dict["input_ids"]
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
>>> doc_scores = tf.squeeze(
...     tf.matmul(
...         tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
...     ),
...     axis=1,
... )
>>> # 3. Forward to generator
>>> outputs = model(
...     inputs=None,
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
...     decoder_input_ids=input_dict["labels"],
... )

>>> # or directly generate
>>> generated = model.generate(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)

generate

< >

( input_ids: TFModelInputType | None = None attention_mask: tf.Tensor | None = None context_input_ids = None context_attention_mask = None doc_scores = None do_deduplication = None num_return_sequences = None num_beams = None n_docs = None **model_kwargs ) tf.Tensor of shape (batch_size * num_return_sequences, sequence_length)

参数

  • input_ids (tf.Tensor of shape (batch_size, sequence_length), optional) — 用作生成的提示的序列。 如果未传递 input_ids,则必须提供 context_input_ids
  • attention_mask (tf.Tensor of shape (batch_size, sequence_length), optional) — 用于避免在填充标记索引上执行注意力的掩码。 在 [0, 1] 中选择的掩码值: - 1 表示未被掩盖的标记, - 0 表示被掩盖的标记。 什么是注意力掩码?
  • context_input_ids (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。
  • context_attention_mask (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。 如果模型没有用 retriever 初始化,或者未给出 input_ids,则必须将 context_input_idscontext_attention_mask 提供给前向传递。 它们由 __call__() 返回。
  • doc_scores (tf.Tensor of shape (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。 如果模型没有用 retriever 初始化,或者未给出 input_ids,则必须将 doc_scores 提供给前向传递。 doc_scores__call__() 返回。
  • do_deduplication (bool, optional) — 是否对给定输入的来自不同上下文文档的生成结果进行去重。 如果在分布式后端进行训练时使用,则必须将其设置为 False
  • num_return_sequences(int, optional, defaults to 1) — 批次中每个元素独立计算的返回序列的数量。 请注意,这不是我们传递给 generator[generate()](/docs/transformers/v4.50.0/en/main_classes/text_generation#transformers.GenerationMixin.generate) 函数的值,我们在其中将 num_return_sequences 设置为 num_beams
  • num_beams (int, optional, defaults to 1) — 集束搜索的束数量。 1 表示不使用集束搜索。
  • n_docs (int, optional, defaults to config.n_docs) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • kwargs (Dict[str, Any], optional) — 其他 kwargs 将传递给 generate()

返回:

tf.Tensor of shape (batch_size * num_return_sequences, sequence_length)

生成的序列。 第二个维度(序列长度)要么等于 max_length,要么更短,如果所有批次由于 eos_token_id 而提前完成。

实现 RAG 序列“彻底”解码。 阅读 generate()` 文档,以获取有关如何设置其他 generate 输入参数的更多信息

TFRagTokenForGeneration

class transformers.TFRagTokenForGeneration

< >

( config: Optional[PretrainedConfig] = None question_encoder: Optional[TFPreTrainedModel] = None generator: Optional[TFPreTrainedModel] = None retriever: Optional[RagRetriever] = None **kwargs )

参数

  • config (RagConfig) — 模型的配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,仅加载配置。请查看 from_pretrained() 方法来加载模型权重。
  • question_encoder (TFPreTrainedModel) — 与 retriever 封装的 faiss 索引兼容的编码器模型。
  • generator (TFPreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
  • retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。

The TFRagTokenForGeneration forward 方法,重写了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

TF RAG-token 模型实现。它在前向传播中执行 RAG-token 特定的边缘化 (marginalization)。

RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在正向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。

问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,生成器可以是任何seq2seq模型,最好是 TFBartForConditionalGeneration

该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中——有关更多详细信息,请参见示例。该模型兼容任何自编码模型作为 question_encoder,以及任何带有语言模型头的 seq2seq 模型作为 generator。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder,使用 TFBartForConditionalGeneration 作为 generator

此模型继承自 TFPreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。

此模型也是 TensorFlow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与一般用法和行为相关的所有事项。

该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。

call

< >

( input_ids: TFModelInputType | None = None attention_mask: np.ndarray | tf.Tensor | None = None decoder_input_ids: np.ndarray | tf.Tensor | None = None decoder_attention_mask: np.ndarray | tf.Tensor | None = None encoder_outputs: np.ndarray | tf.Tensor | None = None past_key_values: Tuple[Tuple[Union[np.ndarray, tf.Tensor]]] | None = None doc_scores: np.ndarray | tf.Tensor | None = None context_input_ids: np.ndarray | tf.Tensor | None = None context_attention_mask: np.ndarray | tf.Tensor | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None output_retrieved: bool | None = None n_docs: int | None = None do_marginalize: bool | None = None labels: np.ndarray | tf.Tensor | None = None reduce_loss: bool | None = None return_dict: bool | None = None training: bool = False **kwargs ) transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput or tuple(tf.Tensor)

参数

  • input_ids (tf.Tensor,形状为 (batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定要使用的生成器,并且还指定兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。
  • attention_mask (tf.Tensor,形状为 (batch_size, sequence_length), 可选) — 用于避免在 padding token 索引上执行 attention 的掩码。 掩码值在 [0, 1] 中选择:

    • 1 表示 未被掩盖 的 tokens,
    • 0 表示 被掩盖 的 tokens。

    什么是 attention 掩码?

  • encoder_outputs (tuple(tuple(tf.Tensor), 可选) — 由以下元素组成的元组 (generator_enc_last_hidden_state, 可选: generator_enc_hidden_states, 可选: generator_enc_attentions)。 形状为 (batch_size, n_docs * sequence_length, hidden_size)generator_enc_last_hidden_state 是生成器编码器最后一层输出端的 hidden-states 序列。

    被 (TFRagModel) 模型在解码期间使用。

  • decoder_input_ids (tf.Tensor,形状为 (batch_size, target_sequence_length), 可选) — 为生成任务提供。 默认为 None,根据您正在 RAG 实例中使用的生成器模型的说明构建。
  • decoder_attention_mask (torch.BoolTensor,形状为 (batch_size, target_sequence_length), 可选) — 默认行为:生成一个 tensor,该 tensor 忽略 decoder_input_ids 中的 pad tokens。默认情况下,也将使用因果掩码 (causal mask)。
  • past_key_values (tuple(tuple(tf.Tensor))) — 由两个元素组成的元组:RAG 模型的 encoder_outputs (参见 encoder_outputs) 和底层生成器的 past_key_values。 可用于加速解码。 past_key_values 在解码期间在 (RagTokenForGeneration) 模型中使用。
  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档 embeddings (参见 retrieved_doc_embeds) 和 question_encoder_last_hidden_state 之间的分数。 如果模型没有使用 retriever 初始化,则必须将 doc_scores 提供给前向传播。 doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,有关更多信息,请参见示例。
  • context_input_ids (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。

    如果模型没有使用 retriever 初始化,则必须将 `context_input_ids 提供给前向传播。 context_input_ids__call__() 返回。 context_attention_mask (tf.Tensor,形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回): 由检索器从检索到的文档和问题编码器 input_ids 经过后处理的 attention 掩码。

    如果模型没有使用 retriever 初始化,则必须将 context_attention_mask 提供给前向传播。 context_attention_mask__call__() 返回。

  • use_cache (bool, 可选, 默认为 True) — 如果设置为 True,则返回 past_key_values 键值状态,并且可以用于加速解码 (参见 past_key_values)。
  • output_attentions (bool, 可选) — 是否返回所有 attention 层的 attentions tensors。 有关更多详细信息,请参见返回的 tensors 下的 attentions
  • output_hidden_states (bool, 可选) — 是否返回所有层的 hidden states。 有关更多详细信息,请参见返回的 tensors 下的 hidden_states
  • output_retrieved(bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。 有关更多详细信息,请参见返回的 tensors。
  • return_dict (bool, 可选) — 是否返回 TFRetrievAugLMOutput 而不是普通元组。
  • n_docs (int, 可选, 默认为 `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • do_marginalize (bool, 可选) — 如果为 True,则通过使用 torch.nn.functional.log_softmax 对所有文档进行边缘化 logits。
  • labels (tf.Tensornp.ndarray,形状为 (batch_size, sequence_length), 可选) — 用于根据 Rag-Token 模型公式计算交叉熵分类损失的标签。 有关 Rag-Token 公式化的详细信息,请参见 https://arxiv.org/pdf/2005.11401.pdf 第 2.1 节。 索引应在 [0, ..., config.vocab_size - 1] 中。
  • reduce_loss (bool, optional) — 仅当传递 labels 时相关。如果为 True,则使用 tf.Tensor.sum 操作来减少 NLL 损失。
  • kwargs (Dict[str, any], optional, defaults to {}) — 旧版字典,这是必需的,以便模型可以使用 generate() 函数。

返回:

transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutputtuple(tf.Tensor)

一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutputtf.Tensor 的元组(如果传递了 return_dict=False 或当 config.return_dict=False 时),其中包含各种元素,具体取决于配置 (RagConfig) 和输入。

  • loss (tf.Tensor of shape (1,), optional, returned when labels is provided) — 语言建模损失。

  • logits (tf.Tensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。 该分数可能针对每个词汇标记在所有文档中进行边缘化处理。

  • past_key_values (List[tf.Tensor]可选,当传递 use_cache=True 或当 config.use_cache=True 时返回) — 长度为 config.n_layerstf.Tensor 列表,其中每个张量的形状为 (2, batch_size, num_heads, sequence_length, embed_size_per_head))。

    包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见 past_key_values 输入)。

  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

  • retrieved_doc_embeds (tf.Tensor,形状为 (batch_size, config.n_docs, hidden_size)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与 question_encoder_last_hidden_state 一起用于计算 doc_scores

  • retrieved_doc_ids (tf.Tensor (int32) of shape (batch_size, config.n_docs), optional, returned when output_retrieved=True) — 检索器检索到的嵌入文档的索引。

  • context_input_ids (tf.Tensor(int32) of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。

  • context_attention_mask (tf.Tensor (int32) of shape (batch_size * config.n_docs, config.max_combined_length), optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的注意力掩码。

  • question_encoder_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。

  • question_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_enc_last_hidden_state (tf.Tensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。

  • generator_enc_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

  • generator_dec_hidden_states (tuple(tf.Tensor)可选,当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入的输出,另一个用于每层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。

  • generator_dec_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=True 或当 config.output_attentions=True 时返回) — tf.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)

    生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。

The TFRagTokenForGeneration forward 方法,重写了 __call__ 特殊方法。

尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module 实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。

示例

>>> import tensorflow as tf
>>> from transformers import AutoTokenizer, RagRetriever, TFRagTokenForGeneration

>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever, from_pt=True)

>>> input_dict = tokenizer.prepare_seq2seq_batch(
...     "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> outputs = model(input_dict, output_retrieved=True)

>>> # or use retriever separately
>>> # 1. Encode
>>> input_ids = input_dict["input_ids"]
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
>>> doc_scores = tf.squeeze(
...     tf.matmul(
...         tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
...     ),
...     axis=1,
... )
>>> # 3. Forward to generator
>>> outputs = model(
...     inputs=None,
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
...     decoder_input_ids=input_dict["labels"],
... )

>>> # or directly generate
>>> generated = model.generate(
...     context_input_ids=docs_dict["context_input_ids"],
...     context_attention_mask=docs_dict["context_attention_mask"],
...     doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)

generate

< >

( input_ids: TFModelInputType | None = None attention_mask: tf.Tensor | None = None context_input_ids = None context_attention_mask = None doc_scores = None n_docs = None generation_config = None logits_processor = [] **kwargs ) tf.Tensor of shape (batch_size * num_return_sequences, sequence_length)

参数

  • input_ids (tf.Tensor of shape (batch_size, sequence_length), optional) — 用作生成的提示的序列。如果未传递 input_ids,则必须提供 context_input_ids
  • attention_mask (tf.Tensor of shape (batch_size, sequence_length), optional) — 用于避免在 padding 标记索引上执行注意力的掩码。在 [0, 1] 中选择的掩码值:

    • 1 表示未被掩盖的标记,
    • 0 表示被掩盖的标记。

    什么是注意力掩码?

  • context_input_ids (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, 当 output_retrieved=True 时返回) — 通过检索器从检索到的文档和问题编码器 input_ids 进行后处理的输入 ID。

    如果模型未通过 retriever 初始化,则必须将 context_input_ids 提供给前向传递。 context_input_ids__call__() 返回。

  • context_attention_mask (tf.Tensor of shape (batch_size * config.n_docs, config.max_combined_length), optional, 当 output_retrieved=True 时返回) — 通过检索器从检索到的文档和问题编码器 input_ids 进行后处理的注意力掩码。

    如果模型未通过 retriever 初始化,则必须将 context_input_ids 提供给前向传递。 context_input_ids__call__() 返回。

  • doc_scores (tf.Tensor of shape (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。

    如果模型未通过 retriever 初始化,则必须将 context_input_ids 提供给前向传递。 context_input_ids__call__() 返回。

  • n_docs (int, optional, 默认为 config.n_docs) — 要检索的文档数量和/或要为其生成答案的文档数量。
  • generation_config (~generation.GenerationConfig, optional) — 要用作生成调用的基本参数化的生成配置。传递给 generate 的与 generation_config 的属性匹配的 **kwargs 将覆盖它们。如果未提供 generation_config,将使用默认值,该默认值具有以下加载优先级:1) 来自 generation_config.json 模型文件(如果存在);2) 来自模型配置。请注意,未指定的参数将继承 GenerationConfig 的默认值,应检查其文档以参数化生成。
  • logits_processor (TFLogitsProcessorList, optional) — 补充从参数和模型配置构建的默认 logits 处理器的自定义 logits 处理器。如果传递的 logits 处理器已经使用参数或模型配置创建,则会抛出错误。
  • kwargs (Dict[str, Any], optional) — generate_config 的临时参数化和/或将转发到模型的 forward 函数的其他特定于模型的 kwargs。

返回:

tf.Tensor of shape (batch_size * num_return_sequences, sequence_length)

生成的序列。第二个维度(sequence_length)要么等于 max_length,要么如果所有批次由于 eos_token_id 提前完成,则更短。

实现了 TFRAG 令牌解码。

< > 在 GitHub 上更新