Transformers 文档

RAG

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

RAG

PyTorch TensorFlow FlashAttention

概述

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

它基于 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 合著的论文 《用于知识密集型 NLP 任务的检索增强生成》

论文摘要如下:

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

该模型由 ola13 贡献。

使用技巧

检索增强生成 (“RAG”) 模型结合了预训练的密集检索 (DPR) 和 Seq2Seq 模型的强大功能。RAG 模型检索文档,将其传递给 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.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 的 ID。通常设置为 eos_token_id

RagConfig 存储 Rag模型 的配置。配置对象继承自 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 特定输出

class transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput

< >

( loss: typing.Optional[torch.FloatTensor] = None logits: typing.Optional[torch.FloatTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None past_key_values: typing.Optional[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[tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[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=Trueconfig.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 中后处理得到的输入 ID。
  • 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=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 (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器最后一层输出的隐藏状态序列。
  • 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 之后,用于计算交叉注意力头中的加权平均。

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

class transformers.models.rag.modeling_rag.RetrievAugLMOutput

< >

( logits: typing.Optional[torch.FloatTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None past_key_values: typing.Optional[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[tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[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]可选,当传递 use_cache=Trueconfig.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 中后处理得到的输入 ID。
  • 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=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 (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 模型生成器编码器最后一层输出的隐藏状态序列。
  • 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 之后,用于计算交叉注意力头中的加权平均。

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,可选,默认为配置中定义的索引) — 如果指定,则使用此索引而不是使用配置构建的索引。

检索器用于从向量查询中获取文档。它检索文档嵌入以及文档内容,并将其格式化以便与 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 supported index (e.g., Faiss or other index types depending on your setup)
>>> 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"  # 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 结合。

检索

< >

( question_hidden_states: ndarray n_docs: int ) tuple[np.ndarray, np.ndarray, list[dict]]

参数

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

返回

tuple[np.ndarray, np.ndarray, list[dict]]

包含以下对象的元组

  • retrieved_doc_embeds (形状为 (batch_size, n_docs, dim)np.ndarray) — 每个查询检索到的文档的检索嵌入。
  • doc_ids (形状为 (batch_size, n_docs)np.ndarray) — 索引中文档的 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 (PretrainedConfig, optional) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (PreTrainedModel, optional) — 负责将问题编码为用于检索的隐藏状态的模型。
  • generator (PreTrainedModel, optional) — 负责根据检索到的文档生成文本的模型。
  • retriever (RagRetriever, optional) — 负责根据编码问题从知识库中检索文档的组件。

裸露的 Rag 模型输出原始隐藏状态,顶部没有任何特定头部。

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

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

前向

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[tuple[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[tuple[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.RetrievAugLMOutput or tuple(torch.FloatTensor)

参数

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

    什么是输入 ID?

  • attention_mask (形状为 (batch_size, sequence_length)torch.Tensor, 可选) — 避免对填充标记索引执行注意力的掩码。掩码值选择在 [0, 1] 之间:

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

    什么是注意力掩码?

  • 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_input_ids (形状为 (batch_size, target_sequence_length)torch.LongTensor, 可选) — 词汇表中解码器输入序列标记的索引。

    索引可以使用 AutoTokenizer 获取。有关详细信息,请参阅 PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是解码器输入 ID?

  • decoder_attention_mask (形状为 (batch_size, target_sequence_length)torch.BoolTensor, 可选) — 默认行为:生成一个忽略 decoder_input_ids 中填充标记的张量。默认情况下也将使用因果掩码。
  • past_key_values (tuple[tuple[torch.FloatTensor]], 可选) — 预先计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常包括模型在解码上一阶段返回的 past_key_values,当 use_cache=Trueconfig.use_cache=True 时。

    允许两种格式:

    • Cache 实例,请参阅我们的 kv cache 指南
    • 长度为 config.n_layerstuple(torch.FloatTensor) 元组,每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量)。这也被称为旧版缓存格式。

    模型将输出与作为输入提供的缓存格式相同的缓存格式。如果未传递 past_key_values,则将返回旧版缓存格式。

    如果使用 past_key_values,用户可以选择只输入形状为 (batch_size, 1) 的最后 input_ids(那些没有将其过去的键值状态提供给此模型的)而不是形状为 (batch_size, sequence_length) 的所有 input_ids

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor, 可选) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的得分。如果模型未用 retriever 初始化,则必须向正向传递提供 doc_scoresdoc_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_idscontext_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_maskcontext_attention_mask__call__() 返回。
  • use_cache (bool, 可选) — 如果设置为 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, 可选) — 要检索的文档数量。

返回

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 (形状为 (batch_size, sequence_length, config.vocab_size)torch.FloatTensor) — 语言建模头部的预测分数。该分数可能已针对每个词汇标记的所有文档进行边际化。

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(参见 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 (形状为 (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 后处理得到的输入 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=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 (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor, 可选) — 模型生成器编码器最后一层输出的隐藏状态序列。

  • 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 之后,用于计算交叉注意力头部中的加权平均。

RagModel 前向方法,覆盖 __call__ 特殊方法。

虽然前向传递的配方需要在此函数中定义,但之后应调用 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 (PretrainedConfig, 可选) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (PreTrainedModel, 可选) — 负责将问题编码为用于检索的隐藏状态的模型。
  • generator (PreTrainedModel, 可选) — 负责根据检索到的文档生成文本的模型。
  • retriever (RagRetriever, 可选) — 负责根据编码问题从知识库中检索文档的组件。

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

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

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

前向

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[tuple[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[tuple[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) — 词汇表中输入序列标记的索引。RagConfig 用于初始化模型,它还指定要使用的兼容生成器分词器。使用该分词器类来获取索引。

    什么是输入 ID?

  • attention_mask (形状为 (batch_size, sequence_length)torch.Tensor, 可选) — 避免对填充标记索引执行注意力的掩码。掩码值选择在 [0, 1] 之间:

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

    什么是注意力掩码?

  • 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_input_ids (形状为 (batch_size, target_sequence_length)torch.LongTensor, 可选) — 词汇表中解码器输入序列标记的索引。

    索引可以使用 AutoTokenizer 获取。有关详细信息,请参阅 PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是解码器输入 ID?

  • decoder_attention_mask (torch.BoolTensor, 形状为 (batch_size, target_sequence_length), 可选) — 默认行为:生成一个忽略 decoder_input_ids 中填充标记的张量。默认情况下也将使用因果掩码。
  • past_key_values (tuple[tuple[torch.Tensor]], 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常包括模型在解码上一阶段返回的 past_key_values,当 use_cache=Trueconfig.use_cache=True 时。

    允许两种格式:

    • 一个 Cache 实例,请参阅我们的 kv cache 指南
    • 一个 tuple(torch.FloatTensor) 的元组,长度为 config.n_layers,每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量)。这也被称为传统缓存格式。

    模型将输出与输入相同的缓存格式。如果未传递 past_key_values,则将返回传统缓存格式。

    如果使用 past_key_values,用户可以选择仅输入最后一个 input_ids(即未将其过去的键值状态提供给此模型的那些),形状为 (batch_size, 1),而不是所有形状为 (batch_size, sequence_length)input_ids

  • context_input_ids (torch.LongTensor, 形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的输入 ID。如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。
  • context_attention_mask (torch.LongTensor, 形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的注意力掩码。如果模型未用 retriever 初始化,则必须向正向传递提供 context_attention_maskcontext_attention_mask__call__() 返回。
  • doc_scores (torch.FloatTensor, 形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。如果模型未用 retriever 初始化,则必须向正向传递提供 doc_scoresdoc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,更多信息请参见示例。
  • use_cache (bool, 可选) — 如果设置为 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。更多详细信息请参见返回张量。
  • exclude_bos_score (bool, 可选) — 仅当传递 labels 时相关。如果为 True,则在计算损失时忽略 BOS 标记的分数。
  • reduce_loss (bool, 可选) — 仅当传递 labels 时相关。如果为 True,则使用 torch.Tensor.sum 操作来减少 NLL 损失。
  • labels (torch.LongTensor, 形状为 (batch_size, sequence_length), 可选) — 用于计算掩码语言建模损失的标签。索引应为 [0, ..., config.vocab_size] 或 -100(参见 input_ids 文档字符串)。索引设置为 -100 的标记将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size] 中的标记计算。
  • n_docs (int, 可选) — 要检索的文档数量。

返回

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

一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组(如果传递了 return_dict=Falseconfig.return_dict=False),包含取决于配置 (RagConfig) 和输入的不同元素。

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

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

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(参见 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 (形状为 (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 后处理得到的输入 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=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 (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor, 可选) — 模型生成器编码器最后一层输出的隐藏状态序列。

  • 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 之后,用于计算交叉注意力头部中的加权平均。

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

虽然前向传递的配方需要在此函数中定义,但之后应调用 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,
... )

生成

< >

( 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 形状为 (batch_size * num_return_sequences, sequence_length)

参数

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

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

    什么是注意力掩码?

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

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

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

    如果模型未用 retriever 初始化,或者未提供 input_ids,则必须向正向传递提供 doc_scoresdoc_scores__call__() 返回。

  • do_deduplication (bool, 可选) — 是否对给定输入的来自不同上下文文档的生成进行去重。在分布式后端训练时,必须将其设置为 False
  • num_return_sequences(int, 可选, 默认为 1) — 批处理中每个元素独立计算的返回序列的数量。请注意,这与我们传递给 generator[generate()](/docs/transformers/v4.53.3/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,形状为 (batch_size * num_return_sequences, sequence_length)

生成的序列。第二个维度(序列长度)等于 max_length,或者如果所有批次由于 eos_token_id 而提前结束,则更短。

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

RagTokenForGeneration

class 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 (PretrainedConfig, 可选) — 模型配置类,包含模型的所有参数。使用配置文件初始化模型不会加载与模型相关的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (PreTrainedModel, 可选) — 负责将问题编码为隐藏状态以进行检索的模型。
  • generator (PreTrainedModel, 可选) — 负责根据检索到的文档生成文本的模型。
  • retriever (RagRetriever, 可选) — 负责根据编码问题从知识库中检索文档的组件。

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

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

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

前向

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.FloatTensor] = None encoder_outputs: typing.Optional[tuple[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[tuple[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)) — 词汇表中输入序列标记的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定了一个兼容的生成器分词器。使用该分词器类获取索引。

    什么是输入 ID?

  • attention_mask (torch.FloatTensor, 形状为 (batch_size, sequence_length), 可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值选择在 [0, 1] 中:

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

    什么是注意力掩码?

  • 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 (torch.LongTensor, 形状为 (batch_size, target_sequence_length), 可选) — 用于生成任务。默认为 None,根据您与 RAG 实例一起使用的生成器模型的说明构建。
  • decoder_input_ids (torch.LongTensor, 形状为 (batch_size, target_sequence_length), 可选) — 解码器输入序列标记在词汇表中的索引。

    索引可以使用 AutoTokenizer 获取。有关详细信息,请参见 PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是解码器输入 ID?

  • decoder_attention_mask (torch.BoolTensor, 形状为 (batch_size, target_sequence_length), 可选) — 默认行为:生成一个忽略 decoder_input_ids 中填充标记的张量。默认情况下也将使用因果掩码。
  • past_key_values (tuple[tuple[torch.Tensor]], 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常包括模型在解码上一阶段返回的 past_key_values,当 use_cache=Trueconfig.use_cache=True 时。

    允许两种格式:

    • 一个 Cache 实例,请参阅我们的 kv cache 指南
    • 一个 tuple(torch.FloatTensor) 的元组,长度为 config.n_layers,每个元组包含 2 个形状为 (batch_size, num_heads, sequence_length, embed_size_per_head) 的张量)。这也被称为传统缓存格式。

    模型将输出与输入相同的缓存格式。如果未传递 past_key_values,则将返回传统缓存格式。

    如果使用 past_key_values,用户可以选择仅输入最后一个 input_ids(即未将其过去的键值状态提供给此模型的那些),形状为 (batch_size, 1),而不是所有形状为 (batch_size, sequence_length)input_ids

  • context_input_ids (torch.LongTensor, 形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的输入 ID。如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。
  • context_attention_mask (torch.LongTensor, 形状为 (batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的注意力掩码。如果模型未用 retriever 初始化,则必须向正向传递提供 context_attention_maskcontext_attention_mask__call__() 返回。
  • doc_scores (torch.FloatTensor, 形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)和 question_encoder_last_hidden_state 之间的分数。如果模型未用 retriever 初始化,则必须向正向传递提供 doc_scoresdoc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,更多信息请参见示例。
  • use_cache (bool, 可选) — 如果设置为 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。更多详细信息请参见返回张量。
  • do_marginalize (bool, 可选) — 如果为 True,则通过使用 torch.nn.functional.log_softmax 对所有文档的 logits 进行边缘化。
  • reduce_loss (bool, 可选) — 仅当传递 labels 时相关。如果为 True,则使用 torch.Tensor.sum 操作来减少 NLL 损失。
  • labels (torch.LongTensor, 形状为 (batch_size, sequence_length), 可选) — 用于计算掩码语言建模损失的标签。索引应为 [0, ..., config.vocab_size] 或 -100(参见 input_ids 文档字符串)。索引设置为 -100 的标记将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size] 中的标记计算。
  • n_docs (int, 可选) — 要检索的文档数量。

返回

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

一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组(如果传递了 return_dict=Falseconfig.return_dict=False),包含取决于配置 (RagConfig) 和输入的不同元素。

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

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

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(参见 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 (形状为 (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 后处理得到的输入 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=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 (形状为 (batch_size, sequence_length, hidden_size)torch.FloatTensor, 可选) — 模型生成器编码器最后一层输出的隐藏状态序列。

  • 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 之后,用于计算交叉注意力头部中的加权平均。

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

虽然前向传递的配方需要在此函数中定义,但之后应调用 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)

生成

< >

( 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.Optional[typing.Callable[[int, torch.Tensor], list[int]]] = None logits_processor: typing.Optional[transformers.generation.logits_process.LogitsProcessorList] = [] stopping_criteria: typing.Optional[transformers.generation.stopping_criteria.StoppingCriteriaList] = [] **kwargs ) torch.LongTensor of shape (batch_size * num_return_sequences, sequence_length)

参数

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

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

    什么是注意力掩码?

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。

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

返回

torch.LongTensor,形状为 (batch_size * num_return_sequences, sequence_length)

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

实现了 RAG 标记解码。

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 前向方法,覆盖 __call__ 特殊方法。

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

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

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

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

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

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

该模型处于开发阶段,因为它目前仅在 eager-mode 中完全支持,并且可能无法以 SavedModel 格式导出。

调用

< >

( 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.TFRetrievAugLMOutputtuple(tf.Tensor)

参数

  • input_ids (tf.Tensor,形状为 (batch_size, sequence_length)) — 词汇表中输入序列标记的索引。RagConfig 用于初始化模型,它指定要使用的生成器,也指定兼容的生成器分词器。使用该分词器类获取索引。
  • 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 中填充标记的张量。默认情况下也将使用因果掩码。
  • 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_scoresdoc_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_idscontext_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_maskcontext_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.TFRetrievAugLMOutput 或一个 tf.Tensor 元组(如果传递了 return_dict=Falseconfig.return_dict=False)包含根据配置 (RagConfig) 和输入的不同元素。

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

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

    包含解码器预先计算的隐藏状态(注意力块中的键和值),可用于(参见 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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每层输出处的隐藏状态以及初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=Trueconfig.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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

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

    生成器解码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

TFRagModel 前向方法,覆盖 __call__ 特殊方法。

虽然前向传递的配方需要在此函数中定义,但之后应调用 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__ 特殊方法。

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

TF RAG-sequence 模型实现。它在正向传递中执行 RAG-sequence 特定的边际化。

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

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

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

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

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

该模型处于开发阶段,因为它目前仅在 eager-mode 中完全支持,并且可能无法以 SavedModel 格式导出。

调用

< >

( 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 用于初始化模型,它指定要使用的生成器,也指定兼容的生成器分词器。使用该分词器类获取索引。
  • 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 中填充标记的张量。默认情况下也将使用因果掩码。
  • 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,形状为 (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`) — 要检索的文档数量和/或要生成答案的文档数量。
  • exclude_bos_score (bool, 可选) — 仅当传入 labels 时相关。如果为 True,则在计算损失时忽略 BOS 标记的分数。
  • labels (tf.Tensornp.ndarray,形状为 (batch_size, sequence_length)可选) — 根据 Rag-Sequence 模型公式计算交叉熵分类损失的标签。有关 Rag-Sequence 公式化细节,请参见 https://huggingface.co/papers/2005.11401 第 2.1 节。索引应在 [0, ..., config.vocab_size - 1] 范围内。
  • reduce_loss (bool, 可选) — 仅当传入 labels 时相关。如果为 True,则使用 tf.Tensor.sum 操作来减少 NLL 损失。
  • kwargs (dict[str, any], 可选,默认为 {}) — 旧版字典,为了模型能使用 generate() 函数而必需。

返回

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

一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput 或一个 tf.Tensor 元组(如果传入 return_dict=Falseconfig.return_dict=False),包含根据配置 (RagConfig) 和输入的不同元素。

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

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

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

    包含解码器预先计算的隐藏状态(注意力块中的键和值),可用于(参见 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),形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。

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

  • context_attention_mask (tf.Tensor (int32),形状为 (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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每层输出处的隐藏状态以及初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=Trueconfig.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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

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

    生成器解码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

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

虽然前向传递的配方需要在此函数中定义,但之后应调用 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)

生成

< >

( 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 形状为 (batch_size * num_return_sequences, sequence_length)

参数

  • input_ids (tf.Tensor,形状为 (batch_size, sequence_length)可选) — 用作生成提示的序列。如果未传入 input_ids,则必须提供 context_input_ids
  • attention_mask (tf.Tensor,形状为 (batch_size, sequence_length)可选) — 掩码,用于避免在填充标记索引上执行注意力。掩码值选择在 [0, 1]:- 1 表示未被掩码的标记,- 0 表示被掩码的标记。什么是注意力掩码?
  • 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 后处理的注意力掩码。如果模型未用 retriever 初始化或未给出 input_ids,则必须将 context_input_idscontext_attention_mask 提供给前向传递。它们由 __call__() 返回。
  • doc_scores (tf.Tensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(请参见 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.53.3/en/main_classes/text_generation#transformers.GenerationMixin.generate) 函数的值,在该函数中我们将 num_return_sequences 设置为 num_beams
  • num_beams (int, 可选,默认为 1) — Beam 搜索的束数。1 表示没有束搜索。
  • n_docs (int, 可选,默认为 config.n_docs) — 要检索的文档数量和/或要生成答案的文档数量。
  • kwargs (dict[str, Any], 可选) — 附加 kwargs 将传递给 generate()

返回

tf.Tensor 形状为 (batch_size * num_return_sequences, sequence_length)

生成的序列。第二个维度(序列长度)等于 max_length,或者如果所有批次由于 eos_token_id 而提前结束,则更短。

实现 RAG 序列“彻底”解码。有关如何设置其他生成输入参数的更多信息,请阅读 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 索引。

TFRagTokenForGeneration 的前向方法重写了 __call__ 特殊方法。

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

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

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

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

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

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

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

该模型处于开发阶段,因为它目前仅在 eager-mode 中完全支持,并且可能无法以 SavedModel 格式导出。

调用

< >

( 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.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 中填充标记的张量。默认情况下也将使用因果掩码。
  • 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,形状为 (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“) — 要检索的文档数量和/或要生成答案的文档数量。
  • do_marginalize (bool, 可选) — 如果为 True,则使用 torch.nn.functional.log_softmax 对所有文档的 logits 进行边际化处理。
  • labels (tf.Tensornp.ndarray,形状为 (batch_size, sequence_length)可选) — 根据 Rag-Token 模型公式计算交叉熵分类损失的标签。有关 Rag-Token 公式详情,请参阅 https://huggingface.co/papers/2005.11401 第 2.1 节。索引应在 [0, ..., config.vocab_size - 1] 范围内。
  • reduce_loss (bool, 可选) — 仅当传入 labels 时相关。如果为 True,则使用 tf.Tensor.sum 操作减少 NLL 损失。
  • kwargs (dict[str, any], 可选, 默认为 {}) — 旧版字典,为了使模型能够使用 generate() 函数而必需。

返回

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

一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput 或一个 tf.Tensor 元组(如果传入 return_dict=Falseconfig.return_dict=False),包含根据配置 (RagConfig) 和输入的不同元素。

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

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

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

    包含解码器预先计算的隐藏状态(注意力块中的键和值),可用于(参见 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),形状为 (batch_size, config.n_docs)可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。

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

  • context_attention_mask (tf.Tensor (int32),形状为 (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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    问题编码器在每层输出处的隐藏状态以及初始嵌入输出。

  • question_enc_attentions (tuple(tf.Tensor)可选,当传递 output_attentions=Trueconfig.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=Trueconfig.output_hidden_states=True 时返回) — tf.Tensor 元组(一个用于嵌入输出,一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

    生成器编码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

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

    生成器解码器在每层输出处的隐藏状态以及初始嵌入输出。

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

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

TFRagTokenForGeneration 的前向方法重写了 __call__ 特殊方法。

虽然前向传递的配方需要在此函数中定义,但之后应调用 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)

生成

< >

( 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 形状为 (batch_size * num_return_sequences, sequence_length)

参数

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

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

    什么是注意力掩码?

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未用 retriever 初始化,则必须向正向传递提供 context_input_idscontext_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 的默认值,应查阅其文档以参数化生成。
  • logits_processor (TFLogitsProcessorList, 可选) — 补充从参数和模型配置构建的默认 logits 处理器的自定义 logits 处理器。如果传入已使用参数或模型配置创建的 logits 处理器,则会抛出错误。
  • kwargs (dict[str, Any], 可选) — generate_config 的临时参数化和/或将转发到模型 forward 函数的其他模型特定 kwargs。

返回

tf.Tensor 形状为 (batch_size * num_return_sequences, sequence_length)

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

实现 TFRAG token 解码。

< > 在 GitHub 上更新