Transformers 文档
RAG
并获得增强的文档体验
开始使用
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 模型检索文档 (docs),将它们传递给 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) — 检索批次大小,定义为并发发给封装在 RagRetriever 中的 faiss 索引的查询数。 - dataset (
str
, 可选, 默认为"wiki_dpr"
) — HuggingFace Datasets 中索引数据集的数据集标识符(使用datasets.list_datasets()
列出所有可用的数据集和 id)。 - dataset_split (
str
, 可选, 默认为"train"
) — 要加载的dataset
的哪个拆分。 - index_name (
str
, 可选, 默认为"compressed"
) — 与dataset
关联的索引的索引名称。 可以选择"legacy"
、"exact"
和"compressed"
之一。 - index_path (
str
, 可选) — 磁盘上序列化的 faiss 索引的路径。 - passages_path (
str
, 可选) — 与 faiss 索引兼容的文本段落的路径。 如果使用LegacyIndex
,则为必需 - use_dummy_dataset (
bool
, 可选, 默认为False
) — 是否加载由dataset
指定的数据集的“虚拟”变体。 - label_smoothing (
float
, 可选, 默认为 0.0) — 仅当return_loss
设置为True
时相关。 控制损失计算中标签平滑的epsilon
参数值。 如果设置为 0,则不执行标签平滑。 - do_marginalize (
bool
, 可选, 默认为False
) — 如果为True
,则通过使用torch.nn.functional.log_softmax
对所有文档的 logits 进行边缘化。 - reduce_loss (
bool
, 可选, 默认为False
) — 是否使用torch.Tensor.sum
操作来减少 NLL 损失。 - do_deduplication (
bool
, 可选, 默认为True
) — 是否对给定输入的来自不同上下文文档的生成结果进行去重。 如果在分布式后端进行训练时使用,则必须设置为False
。 - exclude_bos_score (
bool
, 可选, 默认为False
) — 在计算损失时是否忽略 BOS token。 - output_retrieved(
bool
, 可选, 默认为False
) — 如果设置为True
,则返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参阅返回的张量。 - use_cache (
bool
, 可选, 默认为True
) — 模型是否应返回上次键/值注意力(并非所有模型都使用)。 - forced_eos_token_id (
int
, 可选) — 当达到max_length
时,强制作为最后一个生成的 token 的 token id。 通常设置为eos_token_id
。
RagConfig 存储 RagModel 的配置。 配置对象继承自 PretrainedConfig,可用于控制模型输出。 有关更多信息,请阅读 PretrainedConfig 中的文档。
from_question_encoder_generator_configs
< source >( question_encoder_config: PretrainedConfig generator_config: PretrainedConfig **kwargs ) → EncoderDecoderConfig
从预训练的编码器模型配置和解码器模型配置实例化 EncoderDecoderConfig(或派生类)。
RagTokenizer
Rag 特定输出
类 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput
< 源码 >( loss: typing.Optional[torch.FloatTensor] = None logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None retrieved_doc_embeds: typing.Optional[torch.FloatTensor] = None retrieved_doc_ids: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None question_encoder_last_hidden_state: typing.Optional[torch.FloatTensor] = None question_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None )
参数
- loss (
torch.FloatTensor
,形状为(1,)
,可选,当提供labels
时返回) — 语言建模损失。 - logits (
torch.FloatTensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。 对于每个词汇表标记,分数可能在所有文档上进行边缘化。 - doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)与question_encoder_last_hidden_state
之间的分数。 - past_key_values (
List[torch.FloatTensor]
,可选,当传递use_cache=True
或当config.use_cache=True
时返回) — 长度为config.n_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 后处理的输入 ids。 - context_attention_mask (
torch.LongTensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
后处理的注意力掩码。 - question_encoder_last_hidden_state (
torch.FloatTensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 模型的问题编码器池化输出的最后一层输出端的隐藏状态序列。 - question_enc_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —torch.FloatTensor
元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为(batch_size, sequence_length, hidden_size)
。问题编码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- question_enc_attentions (
tuple(torch.FloatTensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —torch.FloatTensor
元组 (每一层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_enc_last_hidden_state (
torch.FloatTensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 模型生成器编码器的最后一层输出端的隐藏状态序列。 - generator_enc_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —torch.FloatTensor
元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为(batch_size, sequence_length, hidden_size)
。生成器编码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- generator_enc_attentions (
tuple(torch.FloatTensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —torch.FloatTensor
元组 (每一层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_dec_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —torch.FloatTensor
元组 (对于嵌入的输出一个,对于每一层的输出一个),形状为(batch_size, sequence_length, hidden_size)
。生成器解码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- generator_dec_attentions (
tuple(torch.FloatTensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —torch.FloatTensor
元组 (每一层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_cross_attentions (
tuple(torch.FloatTensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —torch.FloatTensor
元组 (每一层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的交叉注意力权重,在注意力 softmax 之后,用于计算交叉注意力头中的加权平均值。
检索增强边缘化模型输出的基类。
类 transformers.models.rag.modeling_rag.RetrievAugLMOutput
< 源码 >( logits: FloatTensor = None doc_scores: FloatTensor = None past_key_values: typing.Optional[typing.List[torch.FloatTensor]] = None retrieved_doc_embeds: typing.Optional[torch.FloatTensor] = None retrieved_doc_ids: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None question_encoder_last_hidden_state: typing.Optional[torch.FloatTensor] = None question_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None question_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_last_hidden_state: typing.Optional[torch.FloatTensor] = None generator_enc_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_enc_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_dec_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None generator_cross_attentions: typing.Optional[typing.Tuple[torch.FloatTensor, ...]] = None )
参数
- logits (
torch.FloatTensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。 对于每个词汇表标记,分数可能在所有文档上进行边缘化。 - doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)与question_encoder_last_hidden_state
之间的分数。 - past_key_values (
List[torch.FloatTensor]
, optional, returned whenuse_cache=True
is passed or whenconfig.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 经过后处理的输入 ids。 - context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 - question_encoder_last_hidden_state (形状为
(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
, 可选) — 模型问题编码器池化输出的最后一层的输出端的隐藏状态序列。 - question_enc_hidden_states (
tuple(torch.FloatTensor)
, 可选, 当传递output_hidden_states=True
或config.output_hidden_states=True
时返回) — 形状为(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
元组(一个用于嵌入的输出,一个用于每一层的输出)。问题编码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- question_enc_attentions (
tuple(torch.FloatTensor)
, 可选, 当传递output_attentions=True
或config.output_attentions=True
时返回) — 形状为(batch_size, num_heads, sequence_length, sequence_length)
的torch.FloatTensor
元组(每一层一个)。问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_enc_last_hidden_state (形状为
(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
, 可选) — 模型生成器编码器的最后一层的输出端的隐藏状态序列。 - generator_enc_hidden_states (
tuple(torch.FloatTensor)
, 可选, 当传递output_hidden_states=True
或config.output_hidden_states=True
时返回) — 形状为(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
元组(一个用于嵌入的输出,一个用于每一层的输出)。生成器编码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- generator_enc_attentions (
tuple(torch.FloatTensor)
, 可选, 当传递output_attentions=True
或config.output_attentions=True
时返回) — 形状为(batch_size, num_heads, sequence_length, sequence_length)
的torch.FloatTensor
元组(每一层一个)。生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_dec_hidden_states (
tuple(torch.FloatTensor)
, 可选, 当传递output_hidden_states=True
或config.output_hidden_states=True
时返回) — 形状为(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
元组(一个用于嵌入的输出,一个用于每一层的输出)。生成器解码器在每一层输出端的隐藏状态,加上初始嵌入输出。
- generator_dec_attentions (
tuple(torch.FloatTensor)
, 可选, 当传递output_attentions=True
或config.output_attentions=True
时返回) — 形状为(batch_size, num_heads, sequence_length, sequence_length)
的torch.FloatTensor
元组(每一层一个)。生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均值。
- generator_cross_attentions (
tuple(torch.FloatTensor)
, 可选, 当传递output_attentions=True
或config.output_attentions=True
时返回) — 形状为(batch_size, num_heads, sequence_length, sequence_length)
的torch.FloatTensor
元组(每一层一个)。生成器解码器的交叉注意力权重,在注意力 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) — 用于对问题进行分词的分词器。它用于解码问题,然后使用 generator_tokenizer。
- generator_tokenizer (PreTrainedTokenizer) — 用于 RagModel 的生成器部分的分词器。
- index (
Index
, optional, defaults to the one defined by the configuration) — 如果指定,则使用此索引代替使用配置构建的索引
用于从向量查询中获取文档的检索器。它检索文档嵌入以及文档内容,并对其进行格式化以与 RagModel 一起使用。
示例
>>> # To load the default "wiki_dpr" dataset with 21M passages from wikipedia (index name is 'compressed' or 'exact')
>>> from transformers import RagRetriever
>>> retriever = RagRetriever.from_pretrained(
... "facebook/dpr-ctx_encoder-single-nq-base", dataset="wiki_dpr", index_name="compressed"
... )
>>> # To load your own indexed dataset built with the datasets library. More info on how to build the indexed dataset in examples/rag/use_own_knowledge_dataset.py
>>> from transformers import RagRetriever
>>> dataset = (
... ...
... ) # dataset must be a datasets.Datasets object with columns "title", "text" and "embeddings", and it must have a faiss index
>>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", indexed_dataset=dataset)
>>> # To load your own indexed dataset built with the datasets library that was saved on disk. More info in examples/rag/use_own_knowledge_dataset.py
>>> from transformers import RagRetriever
>>> dataset_path = "path/to/my/dataset" # dataset saved via *dataset.save_to_disk(...)*
>>> index_path = "path/to/my/index.faiss" # faiss index saved via *dataset.get_index("embeddings").save(...)*
>>> retriever = RagRetriever.from_pretrained(
... "facebook/dpr-ctx_encoder-single-nq-base",
... index_name="custom",
... passages_path=dataset_path,
... index_path=index_path,
... )
>>> # To load the legacy index built originally for Rag's paper
>>> from transformers import RagRetriever
>>> retriever = RagRetriever.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", index_name="legacy")
检索器初始化函数。它将索引加载到内存中。
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: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )
参数
- config (RagConfig) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重,只会加载配置。查看 from_pretrained() 方法加载模型权重。
- question_encoder (PreTrainedModel) — 与
retriever
封装的 faiss 索引兼容的编码器模型。 - generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
- retriever (RagRetriever) — 检索器类,封装了 faiss 索引,用于查询以获取当前输入的上下文文档。
RagModel forward 方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder
和任何带有语言模型头的 seq2seq 模型作为 generator
。它已经使用 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
进行了测试。
此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.FloatTensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.FloatTensor]]] = None doc_scores: typing.Optional[torch.FloatTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None n_docs: typing.Optional[int] = None ) → transformers.models.rag.modeling_rag.RetrievAugLMOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor
,形状为(batch_size, sequence_length)
) — 词汇表中输入序列标记的索引。RagConfig,用于初始化模型,指定要使用的生成器,它还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。 - attention_mask (
torch.Tensor
,形状为(batch_size, sequence_length)
,可选) — 避免在 padding 标记索引上执行 attention 的 Mask。在[0, 1]
中选择的 Mask 值:- 1 表示未被 Mask 的标记,
- 0 表示被 Mask 的标记。
- encoder_outputs (
tuple(tuple(torch.FloatTensor)
, 可选) — 元组由 (generator_enc_last_hidden_state
, 可选:generator_enc_hidden_states
, 可选:generator_enc_attentions
) 组成。形状为(batch_size, n_docs * sequence_length, hidden_size)
的generator_enc_last_hidden_state
是生成器编码器最后一层输出端的 hidden-states 序列。由 (RagModel) 模型在解码期间使用。
- decoder_input_ids (
torch.LongTensor
,形状为(batch_size, target_sequence_length)
,可选) — 为生成任务提供。默认为None
,根据您在 RAG 实例中使用的生成器模型的说明构建。 - decoder_attention_mask (
torch.BoolTensor
,形状为(batch_size, target_sequence_length)
,可选) — 默认行为:生成一个忽略decoder_input_ids
中 pad 标记的 tensor。默认情况下,也将使用因果 mask。 - past_key_values (
tuple(tuple(torch.FloatTensor))
) — 元组由两个元素组成:RAG 模型的encoder_outputs
(参见encoder_outputs
)和底层生成器的past_key_values
。可用于加速解码。在 (RagTokenForGeneration) 模型解码期间使用past_key_values
。 - doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。如果模型没有使用retriever
初始化,则必须将doc_scores
提供给 forward 传递。doc_scores
可以通过question_encoder_last_hidden_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 mask。如果模型没有使用retriever
初始化,则必须将context_attention_mask
提供给 forward 传递。context_attention_mask
由__call__()
返回。 - use_cache (
bool
,可选,默认为True
) — 如果设置为True
,则返回past_key_values
键值状态,并可用于加速解码(参见past_key_values
)。 - output_attentions (
bool
,可选) — 是否返回所有 attention 层的 attentions tensor。有关更多详细信息,请参见返回的 tensor 下的attentions
。 - output_hidden_states (
bool
,可选) — 是否返回所有层的 hidden states。有关更多详细信息,请参见返回的 tensor 下的hidden_states
。 - output_retrieved(
bool
, 可选) — 是否返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。有关更多详细信息,请参见返回的 tensor。 - 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 (
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
列表,其中每个 tensor 的形状为(2, batch_size, num_heads, sequence_length, embed_size_per_head)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
retrieved_doc_embeds (
torch.FloatTensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
torch.LongTensor
,形状为(batch_size, config.n_docs)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。 -
context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的输入 id。 -
context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 -
question_encoder_last_hidden_state (形状为
(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
,可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。 -
question_enc_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=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: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )
参数
- config (RagConfig) — 带有模型所有参数的模型配置类。使用配置文件初始化不会加载与模型关联的权重,仅加载配置。查看 from_pretrained() 方法以加载模型权重。
- question_encoder (PreTrainedModel) — 与由
retriever
封装的 faiss 索引兼容的编码器模型。 - generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
- retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,用于查询以获取当前输入的上下文文档。
RagSequenceForGeneration 的 forward 方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG 序列模型实现。它在前向传播中执行 RAG 序列特定的边缘化。
RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder
和任何带有语言模型头的 seq2seq 模型作为 generator
。它已经使用 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
进行了测试。
此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None exclude_bos_score: typing.Optional[bool] = None reduce_loss: typing.Optional[bool] = None labels: typing.Optional[torch.LongTensor] = None n_docs: typing.Optional[int] = None **kwargs ) → transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (形状为
(batch_size, sequence_length)
的torch.LongTensor
) — 输入序列 token 在词汇表中的索引。RagConfig,用于初始化模型,指定要使用的生成器,还指定兼容的生成器 tokenizer。使用该 tokenizer 类获取索引。 - attention_mask (形状为
(batch_size, sequence_length)
的torch.Tensor
,可选) — 用于避免在 padding token 索引上执行注意力的掩码。在[0, 1]
中选择的掩码值:- 1 表示 未被掩盖 的 token,
- 0 表示 被掩盖 的 token。
- encoder_outputs (
tuple(tuple(torch.FloatTensor)
,可选) — 元组由 (generator_enc_last_hidden_state
,可选:generator_enc_hidden_states
,可选:generator_enc_attentions
) 组成。形状为(batch_size, n_docs * sequence_length, hidden_size)
的generator_enc_last_hidden_state
是生成器编码器最后一层输出的隐藏状态序列。由 (RagModel) 模型在解码期间使用。
- decoder_input_ids (形状为
(batch_size, target_sequence_length)
的torch.LongTensor
,可选) — 为生成任务提供。默认为None
,根据您与 RAG 实例一起使用的生成器模型的说明构建。 - decoder_attention_mask (形状为
(batch_size, target_sequence_length)
的torch.BoolTensor
,可选) — 默认行为:生成一个忽略decoder_input_ids
中的 pad token 的 tensor。默认情况下也将使用因果掩码。 - past_key_values (
tuple(tuple(torch.FloatTensor))
) — 元组由两个元素组成:RAG 模型的encoder_outputs
(参见encoder_outputs
) 和底层生成器的past_key_values
。可用于加速解码。past_key_values
在 (RagTokenForGeneration) 模型解码期间使用。 - doc_scores (形状为
(batch_size, config.n_docs)
的torch.FloatTensor
) — 每个检索到的文档嵌入 (参见retrieved_doc_embeds
) 和question_encoder_last_hidden_state
之间的分数。如果模型未用retriever
初始化,则必须将doc_scores
提供给前向传播。doc_scores
可以通过question_encoder_last_hidden_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
提供给前向传播。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 or tuple(torch.FloatTensor)
一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor
元组 (如果传递了 return_dict=False
或者当 config.return_dict=False
时),包含各种元素,具体取决于配置 (RagConfig) 和输入。
-
loss (形状为
(1,)
的torch.FloatTensor
, 可选, 当提供labels
时返回) — 语言建模损失。 -
logits (
torch.FloatTensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。分数可能是针对每个词汇表标记对所有文档进行边缘化的结果。 -
doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 -
past_key_values (
List[torch.FloatTensor]
,可选,当传递use_cache=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的torch.FloatTensor
列表,其中每个 tensor 的形状为(2, batch_size, num_heads, sequence_length, embed_size_per_head)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
retrieved_doc_embeds (
torch.FloatTensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
torch.LongTensor
,形状为(batch_size, config.n_docs)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。 -
context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的输入 id。 -
context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 -
question_encoder_last_hidden_state (形状为
(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
,可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。 -
question_enc_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=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: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None do_deduplication: typing.Optional[bool] = None num_return_sequences: typing.Optional[int] = None num_beams: typing.Optional[int] = None n_docs: typing.Optional[int] = None **model_kwargs ) → torch.LongTensor
of shape (batch_size * num_return_sequences, sequence_length)
参数
- input_ids (形状为
(batch_size, sequence_length)
的torch.LongTensor
, 可选) — 用作生成的提示的序列。 如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (形状为
(batch_size, sequence_length)
的torch.Tensor
, 可选) — 用于避免对填充标记索引执行注意力的掩码。 在[0, 1]
中选择的掩码值:- 1 表示未被掩盖的标记,
- 0 表示被掩盖的标记。
- context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 预处理后的输入 ID。 - context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
, 可选, 当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
预处理后的注意力掩码。如果模型未用
retriever
初始化,或者未给出input_ids
,则必须将context_input_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.50.0/en/main_classes/text_generation#transformers.GenerationMixin.generate)
函数的值,我们在该函数中将num_return_sequences
设置为num_beams
。 - num_beams (
int
, 可选, 默认为 1) — 束搜索的束数量。 1 表示不进行束搜索。 - n_docs (
int
, 可选, 默认为config.n_docs
) — 检索的文档数量和/或生成答案的文档数量。 - kwargs (
Dict[str, Any]
, 可选) — 其他 kwargs 将传递给 generate()。
返回:
torch.LongTensor
of shape (batch_size * num_return_sequences, sequence_length)
生成的序列。 第二个维度(序列长度)要么等于 max_length
,要么更短,如果所有批次由于 eos_token_id
而提前完成。
实现 RAG 序列“彻底”解码。 阅读 generate()` 文档以获取有关如何设置其他生成输入参数的更多信息。
RagTokenForGeneration
类 transformers.RagTokenForGeneration
< source >( config: typing.Optional[transformers.configuration_utils.PretrainedConfig] = None question_encoder: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None generator: typing.Optional[transformers.modeling_utils.PreTrainedModel] = None retriever: typing.Optional[transformers.models.rag.retrieval_rag.RagRetriever] = None **kwargs )
参数
- config (RagConfig) — 带有模型所有参数的模型配置类。 使用配置文件初始化不会加载与模型关联的权重,仅加载配置。 查看 from_pretrained() 方法来加载模型权重。
- question_encoder (PreTrainedModel) — 与
retriever
封装的 faiss 索引兼容的编码器模型。 - generator (PreTrainedModel) — 在 RAG 架构中用作生成器的 seq2seq 模型。
- retriever (RagRetriever) — 一个检索器类,封装了一个 faiss 索引,该索引被查询以获取当前输入的上下文文档。
RagTokenForGeneration 的前向方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG-token 模型实现。 它在前向传递中执行 RAG-token 特定的边缘化。
RAG 是一个 seq2seq 模型,它封装了两个核心组件:问题编码器和生成器。在前向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。
问题编码器可以是任何自编码模型,最好是 DPRQuestionEncoder,而生成器可以是任何seq2seq 模型,最好是 BartForConditionalGeneration。
该模型可以使用 RagRetriever 初始化以进行端到端生成,也可以与检索器的输出结合使用在多个步骤中---请参阅示例了解更多详情。该模型兼容任何自编码模型作为 question_encoder
和任何带有语言模型头的 seq2seq 模型作为 generator
。它已经使用 DPRQuestionEncoder 作为 question_encoder
和 BartForConditionalGeneration 或 T5ForConditionalGeneration 作为 generator
进行了测试。
此模型继承自 PreTrainedModel。查看超类文档,了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch Module,并参阅 PyTorch 文档,了解与常规用法和行为相关的所有事项。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.FloatTensor] = None encoder_outputs: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None decoder_input_ids: typing.Optional[torch.LongTensor] = None decoder_attention_mask: typing.Optional[torch.BoolTensor] = None past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.Tensor]]] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None output_retrieved: typing.Optional[bool] = None do_marginalize: typing.Optional[bool] = None reduce_loss: typing.Optional[bool] = None labels: typing.Optional[torch.LongTensor] = None n_docs: typing.Optional[int] = None **kwargs ) → transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor
,形状为(batch_size, sequence_length)
) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定了要使用的生成器,并且还指定了兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。 - attention_mask (
torch.Tensor
,形状为(batch_size, sequence_length)
,可选) — Mask,用于避免对 padding token 索引执行 attention。Mask 值在[0, 1]
中选择:- 1 表示 **未被 mask** 的 tokens,
- 0 表示 **已被 mask** 的 tokens。
- encoder_outputs (
tuple(tuple(torch.FloatTensor)
, 可选) — Tuple 由 (generator_enc_last_hidden_state
, 可选:generator_enc_hidden_states
, 可选:generator_enc_attentions
) 组成。形状为(batch_size, n_docs * sequence_length, hidden_size)
的generator_enc_last_hidden_state
是生成器 encoder 最后一层输出的 hidden-states 序列。由 (RagModel) 模型在解码期间使用。
- decoder_input_ids (
torch.LongTensor
,形状为(batch_size, target_sequence_length)
,可选) — 为生成任务提供。 默认值为None
,根据您在 RAG 实例中使用的生成器模型的说明构建。 - decoder_attention_mask (
torch.BoolTensor
,形状为(batch_size, target_sequence_length)
,可选) — 默认行为:生成一个 tensor,该 tensor 忽略decoder_input_ids
中的 pad tokens。因果 mask 也将默认使用。 - past_key_values (
tuple(tuple(torch.FloatTensor))
) — Tuple 由两个元素组成:RAG 模型的encoder_outputs
(参见encoder_outputs
) 和底层生成器的past_key_values
。可用于加速解码。past_key_values
在 (RagTokenForGeneration) 模型解码期间使用。 - doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入 (参见retrieved_doc_embeds
) 和question_encoder_last_hidden_state
之间的分数。如果模型没有使用retriever
初始化,则必须将doc_scores
提供给 forward 传递。doc_scores
可以通过question_encoder_last_hidden_state
和retrieved_doc_embeds
计算,有关更多信息,请参见示例。 - context_input_ids (
torch.LongTensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
后处理的输入 IDs。如果模型没有使用retriever
初始化,则必须将 `context_input_ids
提供给 forward 传递。context_input_ids
由__call__()
返回。 - context_attention_mask (
torch.LongTensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
后处理的 Attention mask。如果模型没有使用retriever
初始化,则必须将context_attention_mask
提供给 forward 传递。context_attention_mask
由__call__()
返回。 - use_cache (
bool
,可选,默认为True
) — 如果设置为True
,则返回past_key_values
键值状态,并且可以用于加速解码 (参见past_key_values
)。 - output_attentions (
bool
,可选) — 是否返回所有 attention 层的 attentions tensors。 有关更多详细信息,请参见返回的 tensors 下的attentions
。 - output_hidden_states (
bool
,可选) — 是否返回所有层的 hidden states。 有关更多详细信息,请参见返回的 tensors 下的hidden_states
。 - output_retrieved(
bool
, 可选) — 是否返回retrieved_doc_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 or tuple(torch.FloatTensor)
一个 transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput 或一个 torch.FloatTensor
元组 (如果传递了 return_dict=False
或者当 config.return_dict=False
时),包含各种元素,具体取决于配置 (RagConfig) 和输入。
-
loss (形状为
(1,)
的torch.FloatTensor
, 可选, 当提供labels
时返回) — 语言建模损失。 -
logits (
torch.FloatTensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。分数可能是针对每个词汇表标记对所有文档进行边缘化的结果。 -
doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 -
past_key_values (
List[torch.FloatTensor]
,可选,当传递use_cache=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的torch.FloatTensor
列表,其中每个 tensor 的形状为(2, batch_size, num_heads, sequence_length, embed_size_per_head)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
retrieved_doc_embeds (
torch.FloatTensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档。与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
torch.LongTensor
,形状为(batch_size, config.n_docs)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入式文档的索引。 -
context_input_ids (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的输入 id。 -
context_attention_mask (形状为
(batch_size * config.n_docs, config.max_combined_length)
的torch.LongTensor
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 -
question_encoder_last_hidden_state (形状为
(batch_size, sequence_length, hidden_size)
的torch.FloatTensor
,可选) — 问题编码器模型池化输出的最后一层输出的隐藏状态序列。 -
question_enc_hidden_states (
tuple(torch.FloatTensor)
,可选,当传递output_hidden_states=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: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None context_input_ids: typing.Optional[torch.LongTensor] = None context_attention_mask: typing.Optional[torch.LongTensor] = None doc_scores: typing.Optional[torch.FloatTensor] = None n_docs: typing.Optional[int] = None generation_config: typing.Optional[transformers.generation.configuration_utils.GenerationConfig] = None prefix_allowed_tokens_fn: typing.Callable[[int, torch.Tensor], typing.List[int]] = None logits_processor: typing.Optional[transformers.generation.logits_process.LogitsProcessorList] = [] stopping_criteria: typing.Optional[transformers.generation.stopping_criteria.StoppingCriteriaList] = [] **kwargs ) → torch.LongTensor
,形状为 (batch_size * num_return_sequences, sequence_length)
参数
- input_ids (
torch.LongTensor
,形状为(batch_size, sequence_length)
,可选) — 用作生成的 prompt 的序列。如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (
torch.Tensor
,形状为(batch_size, sequence_length)
,可选) — Mask,用于避免对 padding token 索引执行 attention。Mask 值在[0, 1]
中选择:- 1 表示 **未被 mask** 的 tokens,
- 0 表示 **已被 mask** 的 tokens。
- context_input_ids (
torch.LongTensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
后处理的输入 IDs。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward 传递。context_input_ids
由__call__()
返回。 - context_attention_mask (
torch.LongTensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器input_ids
后处理的 Attention mask。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward 传递。context_input_ids
由__call__()
返回。 - doc_scores (
torch.FloatTensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入 (参见retrieved_doc_embeds
) 和question_encoder_last_hidden_state
之间的分数。如果模型没有使用
retriever
初始化,则必须将context_input_ids
提供给 forward 传递。context_input_ids
由__call__()
返回。 - n_docs (
int
,可选,默认为config.n_docs
) — 要检索的文档数量和/或要为其生成答案的文档数量。 - generation_config (
~generation.GenerationConfig
, 可选) — 用作生成调用的基本参数化的生成配置。传递给 generate 的与generation_config
属性匹配的**kwargs
将覆盖它们。如果未提供generation_config
,则将使用默认值,其加载优先级如下:1) 来自generation_config.json
模型文件(如果存在);2) 来自模型配置。请注意,未指定的参数将继承 GenerationConfig 的默认值,应查阅其文档以参数化生成。 - prefix_allowed_tokens_fn (
Callable[[int, torch.Tensor], List[int]]
, 可选) — 如果提供,此函数将在每一步将束搜索约束为仅允许的 token。如果未提供,则不应用约束。此函数接受 2 个参数inputs_ids
和批次 IDbatch_id
。它必须返回一个列表,其中包含用于下一步生成的允许 token,这些 token 以先前生成的 tokeninputs_ids
和批次 IDbatch_id
为条件。此参数对于以 prefix 为条件的约束生成很有用,如 Autoregressive Entity Retrieval 中所述。 - logits_processor (
LogitsProcessorList
, 可选) — 补充从参数和模型配置构建的默认 logits 处理器的自定义 logits 处理器。如果传递的 logit 处理器已使用参数或模型配置创建,则会抛出错误。 - stopping_criteria (
StoppingCriteriaList
, 可选) — 补充从参数和模型配置构建的默认停止标准的自定义停止标准。如果传递的停止标准已使用参数或模型配置创建,则会抛出错误。 - kwargs (
Dict[str, Any]
, 可选) —generate_config
的特别参数化和/或将转发到模型的forward
函数的其他模型特定 kwargs。
返回:
torch.LongTensor
of shape (batch_size * num_return_sequences, sequence_length)
生成的序列。第二个维度(sequence_length)要么等于 max_length
,要么如果所有批次由于 eos_token_id
提前完成,则更短。
实现 RAG token 解码。
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) — 用作 RAG 架构中生成器的 seq2seq 模型。
- retriever (RagRetriever) — 一个检索器类,封装了 faiss 索引,该索引被查询以获取当前输入的上下文文档。
TFRagModel
的 forward 方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在正向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。
问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,生成器可以是任何seq2seq模型,最好是 TFBartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中——有关更多详细信息,请参见示例。该模型兼容任何自编码模型作为 question_encoder
,以及任何带有语言模型头的 seq2seq 模型作为 generator
。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder
,使用 TFBartForConditionalGeneration 作为 generator
。
此模型继承自 TFPreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 TensorFlow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与一般用法和行为相关的所有事项。
该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。
call
< 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
or tuple(tf.Tensor)
参数
- input_ids (
tf.Tensor
of shape(batch_size, sequence_length)
) — 词汇表中输入序列 token 的索引。RagConfig,用于初始化模型,指定要使用的生成器,还指定兼容的生成器 tokenizer。 - attention_mask (
tf.Tensor
of shape(batch_size, sequence_length)
, 可选) — 用于避免对 padding token 索引执行 attention 的掩码。掩码值在[0, 1]
中选择:- 1 代表未被掩码的 token,
- 0 代表被掩码的 token。
- encoder_outputs (
tuple(tuple(tf.Tensor)
, 可选) — 元组由 (generator_enc_last_hidden_state
, 可选:generator_enc_hidden_states
, 可选:generator_enc_attentions
) 组成。generator_enc_last_hidden_state
的形状为(batch_size, n_docs * sequence_length, hidden_size)
,是生成器编码器最后一层输出的 hidden-state 序列。由 (TFRagModel) 模型在解码期间使用。
- decoder_input_ids (
tf.Tensor
of shape(batch_size, target_sequence_length)
, 可选) — 为生成任务提供。默认为None
,按照您与 RAG 实例一起使用的生成器模型的说明构建。 - decoder_attention_mask (
torch.BoolTensor
of shape(batch_size, target_sequence_length)
, 可选) — 默认行为:生成一个忽略decoder_input_ids
中 pad token 的 tensor。默认情况下还将使用因果掩码。 - past_key_values (
tuple(tuple(tf.Tensor))
) — 元组由两个元素组成:RAG 模型的encoder_outputs
(参见encoder_outputs
)和底层生成器的past_key_values
。可用于加速解码。past_key_values
在 (RagTokenForGeneration) 模型解码期间使用。 - doc_scores (
tf.Tensor
of shape(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。如果模型未用retriever
初始化,则必须将doc_scores
提供给 forward 传递。doc_scores
可以通过question_encoder_last_hidden_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
) — 要检索的文档数和/或要为其生成答案的文档数。
返回:
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)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
doc_scores (
tf.Tensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(请参阅retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 -
retrieved_doc_embeds (
tf.Tensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
tf.Tensor
,形状为(batch_size, config.n_docs)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档的索引。 -
context_input_ids (
tf.Tensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 从检索到的文档和问题编码器 input_ids 经过检索器后处理的输入 id。 -
context_attention_mask (
tf.Tensor
,形状为(batch_size * config.n_docs, config.max_combined_length)
,可选,当 output_retrieved=True 时返回) — 从检索到的文档和问题编码器input_ids
经过检索器后处理的注意力掩码。 -
question_encoder_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。 -
question_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
question_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_enc_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。 -
generator_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_dec_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_dec_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
TFRagModel
的 forward 方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
示例
>>> from transformers import AutoTokenizer, RagRetriever, TFRagModel
>>> import torch
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base")
>>> retriever = RagRetriever.from_pretrained(
... "facebook/rag-token-base", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagModel.from_pretrained("facebook/rag-token-base", retriever=retriever, from_pt=True)
>>> input_dict = tokenizer.prepare_seq2seq_batch(
... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> input_ids = input_dict["input_ids"]
>>> outputs = model(input_ids)
TFRagSequenceForGeneration
class transformers.TFRagSequenceForGeneration
< 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。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 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
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器input_ids
经过后处理的输入 ID。如果模型没有用
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 context_attention_mask (tf.Tensor
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True): 检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。如果模型没有用
retriever
初始化,则必须将context_attention_mask
提供给前向传递。context_attention_mask
由__call__()
返回。 - use_cache (
bool
, optional, defaults toTrue
) — 如果设置为True
,则返回past_key_values
键值状态,并可用于加速解码(请参阅past_key_values
)。 - output_attentions (
bool
, optional) — 是否返回所有注意力层的注意力张量。 有关更多详细信息,请参阅返回张量下的attentions
。 - output_hidden_states (
bool
, optional) — 是否返回所有层的隐藏状态。 有关更多详细信息,请参阅返回张量下的hidden_states
。 - output_retrieved(
bool
, optional) — 是否返回retrieved_doc_embeds
、retrieved_doc_ids
、context_input_ids
和context_attention_mask
。 有关更多详细信息,请参阅返回张量。 - return_dict (
bool
, optional) — 是否返回TFRetrievAugLMOutput
而不是普通元组。 - n_docs (
int
, optional, defaults to `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。 - exclude_bos_score (
bool
, optional) — 仅在传递labels
时相关。 如果为True
,则在计算损失时忽略 BOS 标记的分数。 - labels (
tf.Tensor
或np.ndarray
of shape(batch_size, sequence_length)
, optional) — 用于根据 Rag-Sequence 模型公式计算交叉熵分类损失的标签。 有关 Rag-Sequence 公式的详细信息,请参阅 https://arxiv.org/pdf/2005.11401.pdf 第 2.1 节。 索引应在[0, ..., config.vocab_size - 1]
中。 - reduce_loss (
bool
, optional) — 仅在传递labels
时相关。 如果为True
,则使用tf.Tensor.sum
操作减少 NLL 损失。 - kwargs (
Dict[str, any]
, optional, defaults to{}
) — 旧版字典,模型需要它才能使用 generate() 函数。
返回:
transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tuple(tf.Tensor)
一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tf.Tensor
的元组(如果传递了 return_dict=False
或当 config.return_dict=False
时),其中包含各种元素,具体取决于配置 (RagConfig) 和输入。
-
loss (
tf.Tensor
of shape(1,)
, optional, returned whenlabels
is provided) — 语言建模损失。 -
logits (
tf.Tensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。 该分数可能针对每个词汇标记在所有文档中进行边缘化处理。 -
past_key_values (
List[tf.Tensor]
,可选,当传递use_cache=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的tf.Tensor
列表,其中每个张量的形状为(2, batch_size, num_heads, sequence_length, embed_size_per_head)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
doc_scores (
tf.Tensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(请参阅retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 -
retrieved_doc_embeds (
tf.Tensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
tf.Tensor
(int32) of shape(batch_size, config.n_docs)
, optional, returned when output_retrieved=True) — 检索器检索到的嵌入文档的索引。 -
context_input_ids (
tf.Tensor
(int32) of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。 -
context_attention_mask (
tf.Tensor
(int32) of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 -
question_encoder_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。 -
question_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
question_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_enc_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。 -
generator_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_dec_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_dec_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
TFRagSequenceForGeneration
前向方法,覆盖了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
示例
>>> from transformers import AutoTokenizer, RagRetriever, TFRagSequenceForGeneration
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq")
>>> retriever = RagRetriever.from_pretrained(
... "facebook/rag-sequence-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagSequenceForGeneration.from_pretrained(
... "facebook/rag-sequence-nq", retriever=retriever, from_pt=True
... )
>>> input_dict = tokenizer.prepare_seq2seq_batch(
... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> outputs = model(input_dict, output_retrieved=True)
>>> # or use retriever separately
>>> # 1. Encode
>>> input_ids = input_dict["input_ids"]
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
>>> doc_scores = tf.squeeze(
... tf.matmul(
... tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
... ),
... axis=1,
... )
>>> # 3. Forward to generator
>>> outputs = model(
... inputs=None,
... context_input_ids=docs_dict["context_input_ids"],
... context_attention_mask=docs_dict["context_attention_mask"],
... doc_scores=doc_scores,
... decoder_input_ids=input_dict["labels"],
... )
>>> # or directly generate
>>> generated = model.generate(
... context_input_ids=docs_dict["context_input_ids"],
... context_attention_mask=docs_dict["context_attention_mask"],
... doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
generate
< 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
of shape (batch_size * num_return_sequences, sequence_length)
参数
- input_ids (
tf.Tensor
of shape(batch_size, sequence_length)
, optional) — 用作生成的提示的序列。 如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (
tf.Tensor
of shape(batch_size, sequence_length)
, optional) — 用于避免在填充标记索引上执行注意力的掩码。 在[0, 1]
中选择的掩码值: - 1 表示未被掩盖的标记, - 0 表示被掩盖的标记。 什么是注意力掩码? - context_input_ids (
tf.Tensor
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。 - context_attention_mask (
tf.Tensor
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 如果模型没有用retriever
初始化,或者未给出input_ids
,则必须将context_input_ids
和context_attention_mask
提供给前向传递。 它们由__call__()
返回。 - doc_scores (
tf.Tensor
of shape(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 如果模型没有用retriever
初始化,或者未给出input_ids
,则必须将doc_scores
提供给前向传递。doc_scores
由__call__()
返回。 - do_deduplication (
bool
, optional) — 是否对给定输入的来自不同上下文文档的生成结果进行去重。 如果在分布式后端进行训练时使用,则必须将其设置为False
。 - num_return_sequences(
int
, optional, defaults to 1) — 批次中每个元素独立计算的返回序列的数量。 请注意,这不是我们传递给generator
的[generate()](/docs/transformers/v4.50.0/en/main_classes/text_generation#transformers.GenerationMixin.generate)
函数的值,我们在其中将num_return_sequences
设置为num_beams
。 - num_beams (
int
, optional, defaults to 1) — 集束搜索的束数量。 1 表示不使用集束搜索。 - n_docs (
int
, optional, defaults toconfig.n_docs
) — 要检索的文档数量和/或要为其生成答案的文档数量。 - kwargs (
Dict[str, Any]
, optional) — 其他 kwargs 将传递给 generate()
返回:
tf.Tensor
of shape (batch_size * num_return_sequences, sequence_length)
生成的序列。 第二个维度(序列长度)要么等于 max_length
,要么更短,如果所有批次由于 eos_token_id
而提前完成。
实现 RAG 序列“彻底”解码。 阅读 generate()` 文档,以获取有关如何设置其他 generate 输入参数的更多信息
TFRagTokenForGeneration
class transformers.TFRagTokenForGeneration
< 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 索引,用于查询以获取当前输入的上下文文档。
The TFRagTokenForGeneration forward 方法,重写了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
TF RAG-token 模型实现。它在前向传播中执行 RAG-token 特定的边缘化 (marginalization)。
RAG 是一个序列到序列模型,它封装了两个核心组件:问题编码器和生成器。在正向传递期间,我们使用问题编码器对输入进行编码,并将其传递给检索器以提取相关的上下文文档。然后将文档添加到输入的前面。这种上下文化的输入被传递给生成器。
问题编码器可以是任何自编码模型,最好是 TFDPRQuestionEncoder,生成器可以是任何seq2seq模型,最好是 TFBartForConditionalGeneration。
该模型可以使用 RagRetriever 进行端到端生成进行初始化,也可以与检索器的输出结合使用在多个步骤中——有关更多详细信息,请参见示例。该模型兼容任何自编码模型作为 question_encoder
,以及任何带有语言模型头的 seq2seq 模型作为 generator
。它已经过测试,使用 TFDPRQuestionEncoder 作为 question_encoder
,使用 TFBartForConditionalGeneration 作为 generator
。
此模型继承自 TFPreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
此模型也是 TensorFlow keras.Model 子类。将其用作常规 TF 2.0 Keras 模型,并参考 TF 2.0 文档以了解与一般用法和行为相关的所有事项。
该模型处于开发状态,因为它现在仅完全支持 eager 模式,并且可能无法以 SavedModel 格式导出。
call
< 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
or tuple(tf.Tensor)
参数
- input_ids (
tf.Tensor
,形状为(batch_size, sequence_length)
) — 词汇表中输入序列 tokens 的索引。RagConfig,用于初始化模型,指定要使用的生成器,并且还指定兼容的生成器 tokenizer。使用该 tokenizer 类来获取索引。 - attention_mask (
tf.Tensor
,形状为(batch_size, sequence_length)
, 可选) — 用于避免在 padding token 索引上执行 attention 的掩码。 掩码值在[0, 1]
中选择:- 1 表示 未被掩盖 的 tokens,
- 0 表示 被掩盖 的 tokens。
- encoder_outputs (
tuple(tuple(tf.Tensor)
, 可选) — 由以下元素组成的元组 (generator_enc_last_hidden_state
, 可选:generator_enc_hidden_states
, 可选:generator_enc_attentions
)。 形状为(batch_size, n_docs * sequence_length, hidden_size)
的generator_enc_last_hidden_state
是生成器编码器最后一层输出端的 hidden-states 序列。被 (TFRagModel) 模型在解码期间使用。
- decoder_input_ids (
tf.Tensor
,形状为(batch_size, target_sequence_length)
, 可选) — 为生成任务提供。 默认为None
,根据您正在 RAG 实例中使用的生成器模型的说明构建。 - decoder_attention_mask (
torch.BoolTensor
,形状为(batch_size, target_sequence_length)
, 可选) — 默认行为:生成一个 tensor,该 tensor 忽略decoder_input_ids
中的 pad tokens。默认情况下,也将使用因果掩码 (causal mask)。 - past_key_values (
tuple(tuple(tf.Tensor))
) — 由两个元素组成的元组:RAG 模型的encoder_outputs
(参见encoder_outputs
) 和底层生成器的past_key_values
。 可用于加速解码。past_key_values
在解码期间在 (RagTokenForGeneration) 模型中使用。 - doc_scores (
tf.Tensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档 embeddings (参见retrieved_doc_embeds
) 和question_encoder_last_hidden_state
之间的分数。 如果模型没有使用retriever
初始化,则必须将doc_scores
提供给前向传播。doc_scores
可以通过question_encoder_last_hidden_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 层的 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
而不是普通元组。 - n_docs (
int
, 可选, 默认为 `config.n_docs“) — 要检索的文档数量和/或要为其生成答案的文档数量。 - do_marginalize (
bool
, 可选) — 如果为True
,则通过使用torch.nn.functional.log_softmax
对所有文档进行边缘化 logits。 - labels (
tf.Tensor
或np.ndarray
,形状为(batch_size, sequence_length)
, 可选) — 用于根据 Rag-Token 模型公式计算交叉熵分类损失的标签。 有关 Rag-Token 公式化的详细信息,请参见 https://arxiv.org/pdf/2005.11401.pdf 第 2.1 节。 索引应在[0, ..., config.vocab_size - 1]
中。 - reduce_loss (
bool
, optional) — 仅当传递labels
时相关。如果为True
,则使用tf.Tensor.sum
操作来减少 NLL 损失。 - kwargs (
Dict[str, any]
, optional, defaults to{}
) — 旧版字典,这是必需的,以便模型可以使用 generate() 函数。
返回:
transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tuple(tf.Tensor)
一个 transformers.models.rag.modeling_tf_rag.TFRetrievAugLMMarginOutput
或 tf.Tensor
的元组(如果传递了 return_dict=False
或当 config.return_dict=False
时),其中包含各种元素,具体取决于配置 (RagConfig) 和输入。
-
loss (
tf.Tensor
of shape(1,)
, optional, returned whenlabels
is provided) — 语言建模损失。 -
logits (
tf.Tensor
,形状为(batch_size, sequence_length, config.vocab_size)
) — 语言建模头的预测分数。 该分数可能针对每个词汇标记在所有文档中进行边缘化处理。 -
past_key_values (
List[tf.Tensor]
,可选,当传递use_cache=True
或当config.use_cache=True
时返回) — 长度为config.n_layers
的tf.Tensor
列表,其中每个张量的形状为(2, batch_size, num_heads, sequence_length, embed_size_per_head)
)。包含解码器的预计算 hidden-states(attention 块中的键和值),可用于加速顺序解码(参见
past_key_values
输入)。 -
doc_scores (
tf.Tensor
,形状为(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(请参阅retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。 -
retrieved_doc_embeds (
tf.Tensor
,形状为(batch_size, config.n_docs, hidden_size)
,可选,当 output_retrieved=True 时返回) — 检索器检索到的嵌入文档。 与question_encoder_last_hidden_state
一起用于计算doc_scores
。 -
retrieved_doc_ids (
tf.Tensor
(int32) of shape(batch_size, config.n_docs)
, optional, returned when output_retrieved=True) — 检索器检索到的嵌入文档的索引。 -
context_input_ids (
tf.Tensor
(int32) of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器 input_ids 经过后处理的输入 ID。 -
context_attention_mask (
tf.Tensor
(int32) of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, returned when output_retrieved=True) — 检索器从检索到的文档和问题编码器input_ids
经过后处理的注意力掩码。 -
question_encoder_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 问题编码器模型池化输出的最后一层输出处的隐藏状态序列。 -
question_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。问题编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
question_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。问题编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_enc_last_hidden_state (
tf.Tensor
,形状为(batch_size, sequence_length, hidden_size)
,可选) — 模型生成器编码器的最后一层输出处的隐藏状态序列。 -
generator_enc_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器编码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_enc_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器编码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
-
generator_dec_hidden_states (
tuple(tf.Tensor)
,可选,当传递output_hidden_states=True
或当config.output_hidden_states=True
时返回) —tf.Tensor
元组(一个用于嵌入的输出,另一个用于每层的输出),形状为(batch_size, sequence_length, hidden_size)
。生成器解码器在每一层输出的隐藏状态,加上初始嵌入输出。
-
generator_dec_attentions (
tuple(tf.Tensor)
,可选,当传递output_attentions=True
或当config.output_attentions=True
时返回) —tf.Tensor
元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)
。生成器解码器的注意力权重,在注意力 softmax 之后,用于计算自注意力头中的加权平均。
The TFRagTokenForGeneration forward 方法,重写了 __call__
特殊方法。
尽管 forward 传递的步骤需要在该函数中定义,但应该在之后调用 Module
实例而不是此函数,因为前者负责运行预处理和后处理步骤,而后者会默默地忽略它们。
示例
>>> import tensorflow as tf
>>> from transformers import AutoTokenizer, RagRetriever, TFRagTokenForGeneration
>>> tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-nq")
>>> retriever = RagRetriever.from_pretrained(
... "facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True
... )
>>> # initialize with RagRetriever to do everything in one forward call
>>> model = TFRagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever, from_pt=True)
>>> input_dict = tokenizer.prepare_seq2seq_batch(
... "How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="tf"
... )
>>> outputs = model(input_dict, output_retrieved=True)
>>> # or use retriever separately
>>> # 1. Encode
>>> input_ids = input_dict["input_ids"]
>>> question_hidden_states = model.question_encoder(input_ids)[0]
>>> # 2. Retrieve
>>> docs_dict = retriever(input_ids.numpy(), question_hidden_states.numpy(), return_tensors="tf")
>>> doc_scores = tf.squeeze(
... tf.matmul(
... tf.expand_dims(question_hidden_states, axis=1), docs_dict["retrieved_doc_embeds"], transpose_b=True
... ),
... axis=1,
... )
>>> # 3. Forward to generator
>>> outputs = model(
... inputs=None,
... context_input_ids=docs_dict["context_input_ids"],
... context_attention_mask=docs_dict["context_attention_mask"],
... doc_scores=doc_scores,
... decoder_input_ids=input_dict["labels"],
... )
>>> # or directly generate
>>> generated = model.generate(
... context_input_ids=docs_dict["context_input_ids"],
... context_attention_mask=docs_dict["context_attention_mask"],
... doc_scores=doc_scores,
... )
>>> generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
generate
< 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 ) → tf.Tensor
of shape (batch_size * num_return_sequences, sequence_length)
参数
- input_ids (
tf.Tensor
of shape(batch_size, sequence_length)
, optional) — 用作生成的提示的序列。如果未传递input_ids
,则必须提供context_input_ids
。 - attention_mask (
tf.Tensor
of shape(batch_size, sequence_length)
, optional) — 用于避免在 padding 标记索引上执行注意力的掩码。在[0, 1]
中选择的掩码值:- 1 表示未被掩盖的标记,
- 0 表示被掩盖的标记。
- context_input_ids (
tf.Tensor
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, 当 output_retrieved=True 时返回) — 通过检索器从检索到的文档和问题编码器input_ids
进行后处理的输入 ID。如果模型未通过
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - context_attention_mask (
tf.Tensor
of shape(batch_size * config.n_docs, config.max_combined_length)
, optional, 当 output_retrieved=True 时返回) — 通过检索器从检索到的文档和问题编码器input_ids
进行后处理的注意力掩码。如果模型未通过
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - doc_scores (
tf.Tensor
of shape(batch_size, config.n_docs)
) — 每个检索到的文档嵌入(参见retrieved_doc_embeds
)和question_encoder_last_hidden_state
之间的分数。如果模型未通过
retriever
初始化,则必须将context_input_ids
提供给前向传递。context_input_ids
由__call__()
返回。 - n_docs (
int
, optional, 默认为config.n_docs
) — 要检索的文档数量和/或要为其生成答案的文档数量。 - generation_config (
~generation.GenerationConfig
, optional) — 要用作生成调用的基本参数化的生成配置。传递给 generate 的与generation_config
的属性匹配的**kwargs
将覆盖它们。如果未提供generation_config
,将使用默认值,该默认值具有以下加载优先级:1) 来自generation_config.json
模型文件(如果存在);2) 来自模型配置。请注意,未指定的参数将继承 GenerationConfig 的默认值,应检查其文档以参数化生成。 - logits_processor (
TFLogitsProcessorList
, optional) — 补充从参数和模型配置构建的默认 logits 处理器的自定义 logits 处理器。如果传递的 logits 处理器已经使用参数或模型配置创建,则会抛出错误。 - kwargs (
Dict[str, Any]
, optional) —generate_config
的临时参数化和/或将转发到模型的forward
函数的其他特定于模型的 kwargs。
返回:
tf.Tensor
of shape (batch_size * num_return_sequences, sequence_length)
生成的序列。第二个维度(sequence_length)要么等于 max_length
,要么如果所有批次由于 eos_token_id
提前完成,则更短。
实现了 TFRAG 令牌解码。