RAG
概述
检索增强生成(“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 模型检索文档,将其传递给 seq2seq 模型,然后进行边缘化以生成输出。检索器和 seq2seq 模块从预训练模型初始化,并联合微调,从而使检索和生成都能够适应下游任务。
RagConfig
class transformers.RagConfig
< source >( 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) — 检索批次大小,定义为并发发往封装的 faiss 索引的查询数 RagRetriever。 - 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 标记。 - output_retrieved(
bool
, 可选, 默认为False
) — 如果设置为True
,则返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参阅返回的张量。 - use_cache (
bool
, 可选, 默认为True
) — 模型是否应返回最后的键/值注意力 (并非所有模型都使用)。 - forced_eos_token_id (
int
, 可选) — 当达到max_length
时,强制作为最后一个生成的标记的标记 ID。 通常设置为eos_token_id
。
RagConfig 存储 RagModel 的配置。 配置对象继承自 PretrainedConfig,可用于控制模型输出。 阅读 PretrainedConfig 中的文档以获取更多信息。
from_question_encoder_generator_configs
< 源代码 >( question_encoder_config: PretrainedConfig generator_config: PretrainedConfig **kwargs ) → EncoderDecoderConfig
从预训练的编码器模型配置和解码器模型配置实例化 EncoderDecoderConfig(或派生类)。
RagTokenizer
Rag 特定输出
class transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput
< 源代码 >( loss: Optional = None logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: Optional = None retrieved_doc_embeds: Optional = None retrieved_doc_ids: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None question_encoder_last_hidden_state: Optional = None question_enc_hidden_states: Optional = None question_enc_attentions: Optional = None generator_enc_last_hidden_state: Optional = None generator_enc_hidden_states: Optional = None generator_enc_attentions: Optional = None generator_dec_hidden_states: Optional = None generator_dec_attentions: Optional = None generator_cross_attentions: Optional = 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_layers
的torch.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=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 之后,用于计算交叉注意力头中的加权平均值。
检索增强边缘化模型输出的基类。
class transformers.models.rag.modeling_rag.RetrievAugLMOutput
< source >( logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: Optional = None retrieved_doc_embeds: Optional = None retrieved_doc_ids: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None question_encoder_last_hidden_state: Optional = None question_enc_hidden_states: Optional = None question_enc_attentions: Optional = None generator_enc_last_hidden_state: Optional = None generator_enc_hidden_states: Optional = None generator_enc_attentions: Optional = None generator_dec_hidden_states: Optional = None generator_dec_attentions: Optional = None generator_cross_attentions: Optional = 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=True
或config.use_cache=True
时返回) — 长度为config.n_layers
的torch.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=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)
。问题编码器的注意力权重,在 attention 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)
。生成器编码器的注意力权重,在 attention 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)
。生成器解码器的注意力权重,在 attention softmax 之后,用于计算自注意力头中的加权平均。
- generator_cross_attentions (
tuple(torch.FloatTensor)
,可选,当传递output_attentions=True
或config.output_attentions=True
时返回) —torch.FloatTensor
的元组(每一层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的交叉注意力权重,在 attention softmax 之后,用于计算交叉注意力头中的加权平均。
RagRetriever
class transformers.RagRetriever
< source >( 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) — 用于对问题进行分词的 tokenizer。它用于解码问题,然后使用 generator_tokenizer。
- generator_tokenizer (PreTrainedTokenizer) — 用于 RagModel 生成器部分的 tokenizer。
- 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 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")
检索器初始化函数。它将索引加载到内存中。
postprocess_docs
< source >( docs input_strings prefix n_docs return_tensors = None ) → tuple(tensors)
后处理检索到的 docs
并将它们与 input_strings
组合。
retrieve
< source >( 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
的文档。
RagModel
class transformers.RagModel
< source >( config: Optional = None question_encoder: Optional = None generator: Optional = None retriever: Optional = 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 模型,它封装了两个核心组件:问题编码器和生成器。在 forward 传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档预先添加到输入中。这种上下文输入将传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的多步输出结合使用 --- 更多详情请参阅示例。该模型兼容任何自编码模型作为 question_encoder
,以及任何带有语言模型头的 seq2seq 模型作为 generator
。它已经过 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
的测试。
此模型继承自 PreTrainedModel。查看超类文档以了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块,并参阅 PyTorch 文档以了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: Optional = None attention_mask: Optional = None encoder_outputs: Optional = None decoder_input_ids: Optional = None decoder_attention_mask: Optional = None past_key_values: Optional = None doc_scores: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None output_retrieved: Optional = None n_docs: Optional = None ) → transformers.models.rag.modeling_rag.RetrievAugLMOutput or tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor
, 形状为(batch_size, sequence_length)
) — 词汇表中输入序列标记的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。 - attention_mask (
torch.Tensor
, 形状为(batch_size, sequence_length)
, 可选) — 掩码,用于避免在 padding 标记索引上执行 attention。掩码值在[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_attention_mask (
torch.BoolTensor
, 形状为(batch_size, target_sequence_length)
, 可选) — 默认行为:生成一个忽略decoder_input_ids
中 pad 标记的张量。默认情况下,也会使用因果掩码。 - past_key_values (
tuple(tuple(torch.FloatTensor))
) — 元组由两个元素组成: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_state
和retrieved_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 掩码。如果模型未用retriever
初始化,则必须将context_attention_mask
提供给 forward 传递。context_attention_mask
由__call__()
返回。 - use_cache (
bool
, 可选, 默认为True
) — 如果设置为True
,则返回past_key_values
键值状态,并可用于加速解码(请参阅past_key_values
)。 - output_attentions (
bool
, 可选) — 是否返回所有 attention 层的 attention 张量。 有关更多详细信息,请参见返回张量下的attentions
。 - output_hidden_states (
bool
, 可选) — 是否返回所有层的隐藏状态。 有关更多详细信息,请参见返回张量下的hidden_states
。 - output_retrieved(
bool
, 可选) — 是否返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参见返回张量。 - n_docs (
int
, 可选, 默认为 `config.n_docs`) — 检索的文档数量和/或生成答案的文档数量。
返回值
transformers.models.rag.modeling_rag.RetrievAugLMOutput 或 tuple(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_layers
的torch.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 的 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
< source >( config: Optional = None question_encoder: Optional = None generator: Optional = None retriever: Optional = None **kwargs )
参数
- config (RagConfig) — 带有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重,仅加载配置。查看 from_pretrained() 方法来加载模型权重。
- question_encoder (PreTrainedModel) — 与
retriever
封装的 faiss 索引兼容的编码器模型。 - generator (PreTrainedModel) — 用作 RAG 架构中生成器的 seq2seq 模型。
- retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。
RagSequenceForGeneration 的 forward 方法覆盖了 __call__
特殊方法。
尽管 forward 传递的配方需要在该函数中定义,但应该在此之后调用 Module
实例,而不是调用此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG-sequence 模型实现。它在前向传递中执行 RAG-sequence 特定的边缘化操作。
RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在 forward 传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档预先添加到输入中。这种上下文输入将传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的多步输出结合使用 --- 更多详情请参阅示例。该模型兼容任何自编码模型作为 question_encoder
,以及任何带有语言模型头的 seq2seq 模型作为 generator
。它已经过 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
的测试。
此模型继承自 PreTrainedModel。查看超类文档以了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块,并参阅 PyTorch 文档以了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: Optional = None attention_mask: Optional = None encoder_outputs: Optional = None decoder_input_ids: Optional = None decoder_attention_mask: Optional = None past_key_values: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None doc_scores: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None output_retrieved: Optional = None exclude_bos_score: Optional = None reduce_loss: Optional = None labels: Optional = None n_docs: Optional = None **kwargs ) → transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (形状为
(batch_size, sequence_length)
的torch.LongTensor
) — 词汇表中输入序列标记的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。 - 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_attention_mask (形状为
(batch_size, target_sequence_length)
的torch.BoolTensor
, 可选) — 默认行为:生成一个忽略decoder_input_ids
中 pad 标记的张量。默认情况下,也将使用因果掩码。 - 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
提供给 forward 传递。doc_scores
可以通过question_encoder_last_hidden_state
和retrieved_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
提供给 forward 传递。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_ids
和context_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 或 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 (形状为
(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_layers
的torch.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__
特殊方法。
尽管 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
< source >( input_ids: Optional = None attention_mask: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None doc_scores: Optional = None do_deduplication: Optional = None num_return_sequences: Optional = None num_beams: Optional = None n_docs: Optional = None **model_kwargs ) → torch.LongTensor
,形状为 (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_ids
和context_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.45.2/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
< source >( config: Optional = None question_encoder: Optional = None generator: Optional = None retriever: Optional = 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 模型,它封装了两个核心组件:问题编码器和生成器。在 forward 传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档预先添加到输入中。这种上下文输入将传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的多步输出结合使用 --- 更多详情请参阅示例。该模型兼容任何自编码模型作为 question_encoder
,以及任何带有语言模型头的 seq2seq 模型作为 generator
。它已经过 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
的测试。
此模型继承自 PreTrainedModel。查看超类文档以了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块,并参阅 PyTorch 文档以了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: Optional = None attention_mask: Optional = None encoder_outputs: Optional = None decoder_input_ids: Optional = None decoder_attention_mask: Optional = None past_key_values: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None doc_scores: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None output_retrieved: Optional = None do_marginalize: Optional = None reduce_loss: Optional = None labels: Optional = None n_docs: Optional = None **kwargs ) → transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (形状为
(batch_size, sequence_length)
的torch.LongTensor
) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定了兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。 - attention_mask (形状为
(batch_size, sequence_length)
的torch.Tensor
, 可选) — 用于避免在 padding token 索引上执行 attention 的 Mask。Mask 值在[0, 1]
中选择:- 1 表示 未被 Mask 的 tokens,
- 0 表示 被 Mask 的 tokens。
- 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 (形状为
(batch_size, target_sequence_length)
的torch.LongTensor
, 可选) — 为生成任务提供。 默认为None
,根据您在 RAG 实例中使用的生成器模型的说明进行构建。 - decoder_attention_mask (形状为
(batch_size, target_sequence_length)
的torch.BoolTensor
, 可选) — 默认行为:生成一个忽略decoder_input_ids
中的 pad tokens 的 tensor。 默认情况下,也会使用因果 mask。 - 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
) — 每个检索到的文档 embeddings (参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 如果模型没有使用retriever
初始化,则必须将doc_scores
提供给 forward pass。doc_scores
可以通过question_encoder_last_hidden_state
和retrieved_doc_embeds
计算,有关详细信息,请参见示例。 - context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
经过后处理的输入 IDs。 如果模型没有使用retriever
初始化,则必须将 `context_input_ids
提供给 forward pass。context_input_ids
由__call__()
返回。 - context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
经过后处理的 Attention mask。 如果模型没有使用retriever
初始化,则必须将context_attention_mask
提供给 forward pass。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_embeds
,retrieved_doc_ids
,context_input_ids
和context_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 或 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 (形状为
(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_layers
的torch.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 的前向方法,覆盖了 __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
< source >( input_ids: Optional = None attention_mask: Optional = None context_input_ids: Optional = None context_attention_mask: Optional = None doc_scores: Optional = None n_docs: Optional = None generation_config: Optional = None prefix_allowed_tokens_fn: Callable = None logits_processor: Optional = [] stopping_criteria: Optional = [] **kwargs ) → 形状为 (batch_size * num_return_sequences, sequence_length)
的 torch.LongTensor
参数
- input_ids (形状为
(batch_size, sequence_length)
的torch.LongTensor
, 可选) — 用作生成的 prompt 的序列。 如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (形状为
(batch_size, sequence_length)
的torch.Tensor
, 可选) — 用于避免在 padding token 索引上执行 attention 的 Mask。Mask 值在[0, 1]
中选择:- 1 表示 未被 Mask 的 tokens,
- 0 表示 被 Mask 的 tokens。
- context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
经过后处理的输入 IDs。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward pass。context_input_ids
由__call__()
返回。 - context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
经过后处理的 Attention mask。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward pass。context_input_ids
由__call__()
返回。 - doc_scores (形状为
(batch_size, config.n_docs)
的torch.FloatTensor
) — 每个检索到的文档 embeddings (参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward pass。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
和批次 IDbatch_id
。它必须返回一个列表,其中包含基于先前生成的 tokeninputs_ids
和批次 IDbatch_id
的下一个生成步骤的允许 token。此参数对于基于前缀的约束生成非常有用,如 自回归实体检索 中所述。 - logits_processor (
LogitsProcessorList
, 可选) — 自定义 logits 处理器,用于补充从参数和模型配置构建的默认 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 token 解码。
TFRagModel
class transformers.TFRagModel
< source >( 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) — 一个 seq2seq 模型,用作 RAG 架构中的生成器。
- retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。
TFRagModel 的前向方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的配方需要在该函数中定义,但应该在此之后调用 Module
实例,而不是调用此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在前向传播期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档预先添加到输入中。这种上下文化的输入被传递到生成器。
问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 TFBartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中 --- 有关更多详细信息,请参见示例。该模型与任何自编码模型(作为 question_encoder
)和任何带有语言模型头的seq2seq 模型(作为 generator
)兼容。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder
,以及 TFBartForConditionalGeneration 作为 generator
。
此模型继承自 TFPreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪 head 等)。
此模型也是 Tensorflow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与常规用法和行为相关的所有事项。
该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。
call
< source >( 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
或 tuple(tf.Tensor)
参数
- input_ids (
tf.Tensor
,形状为(batch_size, sequence_length)
) — 词汇表中输入序列 token 的索引。RagConfig,用于初始化模型,指定要使用的生成器,还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。 - attention_mask (
tf.Tensor
,形状为(batch_size, sequence_length)
, 可选) — 用于避免在 padding token 索引上执行 attention 的掩码。Mask 值在[0, 1]
中选择:- 1 表示 未被掩盖 的 token,
- 0 表示 被掩盖 的 token。
- encoder_outputs (
tuple(tuple(tf.Tensor)
, 可选) — 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
是生成器编码器最后一层输出处的隐藏状态序列。由 (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 token 的 tensor。默认情况下也会使用因果掩码。 - past_key_values (
tuple(tuple(tf.Tensor))
) — Tuple 由两个元素组成: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_state
和retrieved_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 层的 attention tensor。 有关更多详细信息,请参见返回的 tensor 下的attentions
。 - output_hidden_states (
bool
, 可选) — 是否返回所有层的隐藏状态。 有关更多详细信息,请参见返回的 tensor 下的hidden_states
。 - output_retrieved(
bool
, 可选) — 是否返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参阅返回的张量。 - return_dict (
bool
, 可选) — 是否返回TFRetrievAugLMOutput
而不是纯元组。 - n_docs (
int
, 可选, 默认为 `config.n_docs`) — 要检索的文档数量和/或要为其生成答案的文档数量。
返回值
transformers.models.rag.modeling_tf_rag.TFRetrievAugLMOutput
或 tuple(tf.Tensor)
一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMOutput
或 tf.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_layers
的tf.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=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 的前向方法,覆盖了 __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
< source >( 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。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪 head 等)。
此模型也是 Tensorflow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与常规用法和行为相关的所有事项。
该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。
call
< source >( 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.TFRetrievAugLMMarginOutput
或 tuple(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_state
和retrieved_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_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参阅返回的张量。 - return_dict (
bool
, 可选的) — 是否返回TFRetrievAugLMOutput
而不是普通元组。 - n_docs (
int
, 可选的, 默认为 `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。 - exclude_bos_score (
bool
, 可选的) — 仅当传递labels
时相关。 如果为True
,则在计算损失时忽略 BOS token 的分数。 - labels (
tf.Tensor
或np.ndarray
,形状为(batch_size, sequence_length)
, 可选的) — 用于根据 Rag-Sequence 模型公式计算交叉熵分类损失的标签。 有关 Rag-Sequence 公式的详细信息,请参阅 https://arxiv.org/pdf/2005.11401.pdf 第 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.TFRetrievAugLMMarginOutput
或 tuple(tf.Tensor)
一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tf.Tensor
元组(如果传递 return_dict=False
或当 config.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=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的tf.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=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
< source >( 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)
, 可选的) — 用于避免在填充 token 索引上执行 attention 的掩码。 在[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。 - context_attention_mask (
tf.Tensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
, 可选的, 当 output_retrieved=True 时返回) — 检索器从检索到的文档和问题编码器input_ids
后处理的注意力掩码。 如果模型没有使用retriever
初始化,或者未提供input_ids
,则必须将context_input_ids
和context_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.45.2/en/main_classes/text_generation#transformers.GenerationMixin.generate)
函数的值,我们在其中将num_return_sequences
设置为num_beams
。 - num_beams (
int
, 可选的, 默认为 1) — Beam search 的 beam 数量。 1 表示不使用 beam search。 - 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 序列 “thorough” 解码。 阅读 generate()` 文档以获取有关如何设置其他生成输入参数的更多信息
TFRagTokenForGeneration
class transformers.TFRagTokenForGeneration
< source >( 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) — 与检索器封装的 faiss 索引兼容的编码器模型。
- generator (TFPreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
- retriever (RagRetriever) — 检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。
TFRagTokenForGeneration 的 forward 方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的配方需要在该函数中定义,但应该在此之后调用 Module
实例,而不是调用此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
TF RAG-token 模型实现。它在前向传递中执行 RAG-token 特定的边缘化操作。
RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在前向传播期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档预先添加到输入中。这种上下文化的输入被传递到生成器。
问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 TFBartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中 --- 有关更多详细信息,请参见示例。该模型与任何自编码模型(作为 question_encoder
)和任何带有语言模型头的seq2seq 模型(作为 generator
)兼容。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder
,以及 TFBartForConditionalGeneration 作为 generator
。
此模型继承自 TFPreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪 head 等)。
此模型也是 Tensorflow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与常规用法和行为相关的所有事项。
该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。
call
< source >( 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
或 tuple(tf.Tensor)
参数
- input_ids (形状为
(batch_size, sequence_length)
的tf.Tensor
) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定要使用的生成器,并且还指定兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。 - attention_mask (形状为
(batch_size, sequence_length)
的tf.Tensor
, 可选) — 用于避免在 padding token 索引上执行 attention 的 Mask。Mask 值在[0, 1]
中选择:- 1 表示 未被 Mask 的 tokens,
- 0 表示 已被 Mask 的 tokens。
- encoder_outputs (
tuple(tuple(tf.Tensor)
, 可选) — 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
是生成器编码器最后一层输出的 hidden-states 序列。由 (TFRagModel) 模型在解码期间使用。
- decoder_input_ids (形状为
(batch_size, target_sequence_length)
的tf.Tensor
, 可选) — 为生成任务提供。 默认值为None
,根据您与 RAG 实例一起使用的生成器模型的说明构建。 - decoder_attention_mask (形状为
(batch_size, target_sequence_length)
的torch.BoolTensor
, 可选) — 默认行为:生成一个忽略decoder_input_ids
中的 pad tokens 的 tensor。默认情况下,也将使用因果 Mask。 - past_key_values (
tuple(tuple(tf.Tensor))
) — Tuple 由两个元素组成:RAG 模型的encoder_outputs
(参见encoder_outputs
) 和底层生成器的past_key_values
。可用于加速解码。past_key_values
在解码期间在 (RagTokenForGeneration) 模型中使用。 - doc_scores (形状为
(batch_size, config.n_docs)
的tf.Tensor
) — 每个检索到的文档嵌入 (参见retrieved_doc_embeds
) 和question_encoder_last_hidden_state
之间的分数。如果模型没有用retriever
初始化,则必须将doc_scores
提供给前向传递。doc_scores
可以通过question_encoder_last_hidden_state
和retrieved_doc_embeds
计算得出,有关更多信息,请参见示例。 - context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的tf.Tensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
进行后处理的输入 ID。如果模型没有用
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 context_attention_mask (形状为(batch_size * config.n_docs, config.max_combined_length)
的tf.Tensor
, 可选, 当 output_retrieved=True 时返回): 由 retriever 从检索到的文档和问题编码器input_ids
进行后处理的 Attention mask。如果模型没有用
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_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参见返回的 tensors。 - return_dict (
bool
, 可选) — 是否返回TFRetrievAugLMOutput
而不是普通 tuple。 - n_docs (
int
, 可选, 默认为 “config.n_docs”) — 要检索的文档数量和/或要为其生成答案的文档数量。 - do_marginalize (
bool
, 可选) — 如果为True
,则通过使用torch.nn.functional.log_softmax
将 logits 在所有文档上边缘化。 - labels (形状为
(batch_size, sequence_length)
的tf.Tensor
或np.ndarray
, 可选) — 根据 Rag-Token 模型公式计算交叉熵分类损失的标签。有关 Rag-Token 公式,请参见 https://arxiv.org/pdf/2005.11401.pdf 第 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.TFRetrievAugLMMarginOutput
或 tuple(tf.Tensor)
一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tf.Tensor
元组(如果传递 return_dict=False
或当 config.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=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的tf.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=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 之后,用于计算自注意力头中的加权平均值。
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
< source >( 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 ) → 形状为 (batch_size * num_return_sequences, sequence_length)
的 tf.Tensor
参数
- input_ids (形状为
(batch_size, sequence_length)
的tf.Tensor
, 可选) — 用作生成的 prompt 的序列。 如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (形状为
(batch_size, sequence_length)
的tf.Tensor
, 可选) — 用于避免在 padding token 索引上执行 attention 的 Mask。Mask 值在[0, 1]
中选择:- 1 表示 未被 Mask 的 tokens,
- 0 表示 已被 Mask 的 tokens。
- context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的tf.Tensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
进行后处理的输入 ID。如果模型没有用
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的tf.Tensor
, 可选, 当 output_retrieved=True 时返回) — 由 retriever 从检索到的文档和问题编码器input_ids
进行后处理的 Attention mask。如果模型没有用
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - doc_scores (
tf.Tensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)与question_encoder_last_hidden_state
之间的分数。如果模型未通过
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - n_docs (
int
,可选,默认为config.n_docs
) — 要检索的文档数量和/或为其生成答案的文档数量。 - generation_config (
~generation.GenerationConfig
,可选) — 要用作生成调用的基本参数化的生成配置。传递给与generation_config
属性匹配的 generate 的**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 令牌解码。