Transformers 文档

RAG

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

该模型于 2020 年 5 月 22 日发布在 HF papers 上,并于 2020 年 11 月 16 日贡献给 Hugging Face Transformers。

RAG

FlashAttention

检索增强生成 (RAG) 通过预训练的神经检索器,将预训练语言模型(参数化记忆)与外部数据源(非参数化记忆)的访问结合起来。RAG 在推理过程中获取相关段落并以此为条件进行生成。这通常使答案更具事实性,并允许您通过更改索引而不是重新训练整个模型来更新知识。

您可以在 AI at Meta 组织下找到所有原始 RAG 检查点。

该模型由 ola13 贡献。

点击右侧侧边栏中的 RAG 模型,查看有关如何将 RAG 应用于不同语言任务的更多示例。

以下示例演示了如何使用 AutoModel 生成文本。

自动模型
from transformers import RagRetriever, RagSequenceForGeneration, RagTokenizer


tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained(
    "facebook/rag-sequence-nq", dataset="wiki_dpr", index_name="compressed"
)

model = RagSequenceForGeneration.from_pretrained(
    "facebook/rag-sequence-nq",
    retriever=retriever,
    attn_implementation="flash_attention_2",
    device_map="auto",
)
inputs = tokenizer("How many people live in Paris?", return_tensors="pt").to(model.device)
generated = model.generate(input_ids=inputs["input_ids"])
print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0])

量化通过以较低精度存储权重来减少内存使用。请参阅量化概述以了解受支持的后端。以下示例使用 bitsandbytes 将权重量化为 4 位。

import torch

from transformers import BitsAndBytesConfig, RagRetriever, RagSequenceForGeneration, RagTokenizer


bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
retriever = RagRetriever.from_pretrained(
    "facebook/rag-sequence-nq", dataset="wiki_dpr", index_name="compressed"
)

model = RagSequenceForGeneration.from_pretrained(
    "facebook/rag-sequence-nq",
    retriever=retriever,
    quantization_config=bnb,
    device_map="auto",
)
inputs = tokenizer("How many people live in Paris?", return_tensors="pt").to(model.device)
generated = model.generate(input_ids=inputs["input_ids"])
print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0])

RagConfig

class transformers.RagConfig

< >

(参数列表省略,保持原格式)

参数

  • is_encoder_decoder (bool, 可选, 默认为 True) — 模型是否用作编码器/解码器架构。
  • vocab_size (int, 可选) — 模型的词汇表大小。定义了 input_ids 可以表示的不同标记的数量。
  • prefix (str, 可选) — 在传递给生成器模型之前,添加到每个输入前面的字符串前缀。
  • bos_token_id (int, 可选) — 词汇表中用于流开始(beginning-of-stream)的标记 ID。
  • pad_token_id (int, 可选) — 词汇表中用于填充(padding)的标记 ID。
  • eos_token_id (Union[int, list[int]], 可选) — 词汇表中用于流结束(end-of-stream)的标记 ID。
  • decoder_start_token_id (int, 可选) — 如果编码器-解码器模型以与 bos 不同的标记开始解码,则为该标记的 ID。
  • 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 的拆分(split)。
  • index_name (str, 可选, 默认为 "compressed") — 与 dataset 关联的索引名称。可以在 "legacy""exact""compressed" 之间选择。
  • index_path (str, 可选) — 磁盘上序列化 faiss 索引的路径。
  • passages_path (str, 可选) — 与 faiss 索引兼容的文本段落路径。如果使用 LegacyIndex 则为必需。
  • use_dummy_dataset (bool, 可选, 默认为 False) — 是否加载由 dataset 指定的数据集的“虚拟”(dummy)变体。
  • reduce_loss (bool, 可选, 默认为 False) — 是否使用 torch.Tensor.sum 操作来缩减 NLL 损失。
  • label_smoothing (float, 可选, 默认为 0.0) — 仅在 return_loss 设置为 True 时相关。控制损失计算中标签平滑的 epsilon 参数值。如果设置为 0,则不执行标签平滑。
  • do_deduplication (bool, 可选, 默认为 True) — 是否对给定输入的不同上下文文档生成的文本进行去重。如果在分布式后端训练时使用,必须设置为 False
  • exclude_bos_score (bool, optional, defaults to False) — 是否在计算损失时忽略 BOS token。
  • do_marginalize (bool, optional, defaults to False) — 若为 True,则通过使用 torch.nn.functional.log_softmax 对所有文档的 logits 进行边缘化处理。
  • output_retrieved (bool, optional, defaults to False) — 若设置为 True,将返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。详情请见返回的张量说明。
  • use_cache (bool, optional, defaults to True) — 模型是否应返回最后的 key/values 注意力(并非所有模型都使用)。仅在 config.is_decoder=True 或模型为仅解码器的生成模型时相关。
  • dataset_revision (str, optional,) — 用于检索的 Hugging Face 数据集的修订版本(commit hash、tag 或 branch)。

这是用于存储 RagModel 配置的配置类。它根据指定的参数实例化 Rag 模型,定义了模型架构。使用默认值实例化配置将产生与 类似的配置

配置对象继承自 PreTrainedConfig,可用于控制模型输出。阅读 PreTrainedConfig 的文档以获取更多信息。

from_question_encoder_generator_configs

< >

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

返回

EncoderDecoderConfig

一个配置对象的实例

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

RagTokenizer

class transformers.RagTokenizer

< >

( question_encoder generator )

RAG 特有输出

class transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput

< >

( loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None doc_scores: torch.FloatTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None retrieved_doc_embeds: torch.FloatTensor | None = None retrieved_doc_ids: torch.LongTensor | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | None = None question_encoder_last_hidden_state: torch.FloatTensor | None = None question_enc_hidden_states: tuple[torch.FloatTensor, ...] | None = None question_enc_attentions: tuple[torch.FloatTensor, ...] | None = None generator_enc_last_hidden_state: torch.FloatTensor | None = None generator_enc_hidden_states: tuple[torch.FloatTensor, ...] | None = None generator_enc_attentions: tuple[torch.FloatTensor, ...] | None = None generator_dec_hidden_states: tuple[torch.FloatTensor, ...] | None = None generator_dec_attentions: tuple[torch.FloatTensor, ...] | None = None generator_cross_attentions: tuple[torch.FloatTensor, ...] | None = None )

参数

  • loss (torch.FloatTensor,形状为 (1,)可选,在提供 labels 时返回) — 语言建模损失。
  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能已针对每个词汇 token 在所有文档上进行了边缘化处理。
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的分数。
  • past_key_values (Cache可选,在传递 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。更多详情请参见我们的 kv cache 指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  • loss (torch.FloatTensor,形状为 (1,)可选,在提供 labels 时返回) — 语言建模损失。
  • logits (torch.FloatTensor,形状为 (batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能已针对每个词汇 token 在所有文档上进行了边缘化处理。
  • doc_scores (torch.FloatTensor,形状为 (batch_size, config.n_docs)) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的分数。
  • past_key_values (Cache可选,在传递 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。更多详情请参见我们的 kv cache 指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

class transformers.models.rag.modeling_rag.RetrievAugLMOutput

< >

( logits: torch.FloatTensor | None = None doc_scores: torch.FloatTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None retrieved_doc_embeds: torch.FloatTensor | None = None retrieved_doc_ids: torch.LongTensor | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | None = None question_encoder_last_hidden_state: torch.FloatTensor | None = None question_enc_hidden_states: tuple[torch.FloatTensor, ...] | None = None question_enc_attentions: tuple[torch.FloatTensor, ...] | None = None generator_enc_last_hidden_state: torch.FloatTensor | None = None generator_enc_hidden_states: tuple[torch.FloatTensor, ...] | None = None generator_enc_attentions: tuple[torch.FloatTensor, ...] | None = None generator_dec_hidden_states: tuple[torch.FloatTensor, ...] | None = None generator_dec_attentions: tuple[torch.FloatTensor, ...] | None = None generator_cross_attentions: tuple[torch.FloatTensor, ...] | None = 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 (Cache可选,当传入 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv 缓存指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  • 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 (Cache可选,当传入 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv 缓存指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

RagRetriever

class transformers.RagRetriever

< >

( config question_encoder_tokenizer generator_tokenizer index = None init_retrieval = True )

参数

  • config (RagConfig) — 此检索器所使用的 RAG 模型配置。包含指示要构建哪个 Index 的参数。你可以使用 config.index_name="custom" 加载自己的自定义数据集,或使用 config.index_name="wiki_dpr"(示例)从 datasets 库中使用标准数据集(默认)。
  • question_encoder_tokenizer (PreTrainedTokenizer) — 用于对问题进行分词的分词器。它用于解码问题,然后使用 generator_tokenizer。
  • generator_tokenizer (PreTrainedTokenizer) — 用于 RagModel 生成器部分的分词器。
  • index (Index,可选,默认为配置中定义的索引) — 如果指定,则使用此索引而不是使用配置构建的索引。

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

示例

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

>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-sequence-nq", dataset="wiki_dpr", index_name="compressed"
... )

>>> # To load your own indexed dataset built with the datasets library.
>>> from transformers import RagRetriever

>>> dataset = (
...     ...
... )  # dataset must be a datasets.Datasets object with columns "title", "text" and "embeddings", and it must have a supported index (e.g., Faiss or other index types depending on your setup)
>>> retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", indexed_dataset=dataset)

>>> # To load your own indexed dataset built with the datasets library that was saved on disk.
>>> from transformers import RagRetriever

>>> dataset_path = "path/to/my/dataset"  # dataset saved via *dataset.save_to_disk(...)*
>>> index_path = "path/to/my/index"  # index saved via *dataset.get_index("embeddings").save(...)*
>>> retriever = RagRetriever.from_pretrained(
...     "facebook/rag-sequence-nq",
...     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/rag-sequence-nq", index_name="legacy")

init_retrieval

< >

( )

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

postprocess_docs

< >

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

参数

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

返回

tuple(tensors)

一个元组,包含两个元素:上下文化的 input_ids 和兼容的 attention_mask

对检索到的 docs 进行后处理并将它们与 input_strings 结合。

retrieve

< >

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

参数

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

返回

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

包含以下对象的元组:

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

为指定的 question_hidden_states 检索文档。

RagModel

class transformers.RagModel

< >

( config: transformers.configuration_utils.PreTrainedConfig | None = None question_encoder: transformers.modeling_utils.PreTrainedModel | None = None generator: transformers.modeling_utils.PreTrainedModel | None = None retriever: transformers.models.rag.retrieval_rag.RagRetriever | None = None **kwargs )

参数

  • config (PreTrainedConfig可选) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重,只会加载配置。查看 from_pretrained() 方法以加载模型权重。
  • question_encoder (PreTrainedModel可选) — 负责将问题编码为用于检索的隐藏状态的模型。
  • generator (PreTrainedModel可选) — 负责根据检索到的文档生成文本的模型。
  • retriever (RagRetriever可选) — 负责根据编码后的问题从知识库中检索文档的组件。

原始 Rag 模型,输出未经处理的隐藏状态,顶部没有任何特定的头。

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

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None decoder_input_ids: torch.LongTensor | None = None decoder_attention_mask: torch.BoolTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None doc_scores: torch.FloatTensor | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | 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 **kwargs ) RetrievAugLMOutputtuple(torch.FloatTensor)

参数

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

    什么是输入 ID?

  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — 用于避免对填充 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 (torch.LongTensor,形状为 (batch_size, target_sequence_length)可选) — 为生成任务提供。默认为 None,请根据您与 RAG 实例一起使用的生成器模型的说明进行构建。
  • decoder_input_ids (torch.LongTensor,形状为 (batch_size, target_sequence_length)可选) — 词汇表中解码器输入序列 token 的索引。

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

    什么是解码器输入 ID?

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

    输入仅允许使用 Cache 实例,请参阅我们的 kv 缓存指南。如果未传递 past_key_values,默认将初始化 DynamicCache

    模型将输出与作为输入馈送的缓存格式相同的缓存。

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即那些其键值状态未提供给该模型的 input_ids),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids

  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的得分。如果模型未初始化 retriever,则必须在前向传递中提供 doc_scoresdoc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,更多信息请参阅示例。
  • context_input_ids (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,在 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理得到的输入 ID。如果模型未初始化 retriever,则必须在前向传递中提供 context_input_idscontext_input_ids__call__() 返回。
  • context_attention_mask (形状为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,在 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 后处理得到的注意力掩码。如果模型未初始化 retriever,则必须在前向传递中提供 context_attention_maskcontext_attention_mask__call__() 返回。
  • use_cache (bool, 可选) — 如果设为 True,将返回 past_key_values 键值状态,并可用于加速解码(参见 past_key_values)。
  • output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。更多详细信息请见返回张量中的 attentions
  • output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。更多详细信息请见返回张量中的 hidden_states
  • output_retrieved (bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。更多详细信息请见返回张量。
  • n_docs (int, 可选) — 要检索的文档数量。

返回

RetrievAugLMOutputtuple(torch.FloatTensor)

一个 RetrievAugLMOutput 对象或一个 torch.FloatTensor 元组(如果传递了 return_dict=False 或当 config.return_dict=False 时),根据配置 (RagConfig) 和输入,包含各种元素。

RagModel 的前向传递方法,覆盖了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用 Module 实例而不是这个,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。

  • 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 (Cache,*可选*,当传入 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。欲了解更多细节,请参阅我们的 KV 缓存指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

示例

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

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

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

RagSequenceForGeneration

class transformers.RagSequenceForGeneration

< >

( config: transformers.configuration_utils.PreTrainedConfig | None = None question_encoder: transformers.modeling_utils.PreTrainedModel | None = None generator: transformers.modeling_utils.PreTrainedModel | None = None retriever: transformers.models.rag.retrieval_rag.RagRetriever | None = None **kwargs )

参数

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

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

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

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None encoder_outputs: tuple[tuple[torch.Tensor]] | None = None decoder_input_ids: torch.LongTensor | None = None decoder_attention_mask: torch.BoolTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | None = None doc_scores: torch.FloatTensor | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None output_retrieved: bool | None = None exclude_bos_score: bool | None = None reduce_loss: bool | None = None labels: torch.LongTensor | None = None n_docs: int | None = None **kwargs ) RetrievAugLMMarginOutputtuple(torch.FloatTensor)

参数

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

    什么是输入 ID?

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

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

    什么是注意力掩码?

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

    由 (RagModel) 模型在解码过程中使用。

  • decoder_input_ids (形状为 (batch_size, target_sequence_length)torch.LongTensor可选) — 提供给生成任务使用。默认值为 None,请根据您在 RAG 实例中使用的生成器模型的说明进行构建。
  • decoder_input_ids (shape为 (batch_size, target_sequence_length)torch.LongTensor可选) — 词表中解码器输入序列标记的索引。

    可以使用 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode()PreTrainedTokenizer.call()

    什么是解码器输入 ID?

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

    输入仅允许使用 Cache 实例,详情请参阅我们的 kv 缓存指南。如果不传递 past_key_values,默认将初始化 DynamicCache

    模型将输出与输入相同格式的缓存。

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即那些没有将键值状态提供给此模型的输入),其形状为 (batch_size, unprocessed_length),而不是所有形状为 (batch_size, sequence_length)input_ids

  • context_input_ids (shape为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的输入 ID。如果模型在初始化时没有传入 retriever,则必须在向前传递(forward pass)时提供 context_input_idscontext_input_ids__call__() 返回。
  • context_attention_mask (shape为 (batch_size * config.n_docs, config.max_combined_length)torch.LongTensor可选,当 output_retrieved=True 时返回) — 由检索器从检索到的文档和问题编码器 input_ids 处理后的注意力掩码。如果模型在初始化时没有传入 retriever,则必须在向前传递时提供 context_attention_maskcontext_attention_mask__call__() 返回。
  • doc_scores (shape为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的得分。如果模型在初始化时没有传入 retriever,则必须在向前传递时提供 doc_scoresdoc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算得出,详情请参阅示例。
  • use_cache (bool可选) — 如果设为 True,则返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • output_attentions (bool可选) — 是否返回所有注意力层的注意力张量。详情请参阅返回张量中的 attentions
  • output_hidden_states (bool可选) — 是否返回所有层的隐藏状态。详情请参阅返回张量中的 hidden_states
  • output_retrieved (bool可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。详情请参阅返回的张量。
  • exclude_bos_score (bool可选) — 仅在传入 labels 时相关。如果为 True,则在计算损失时忽略 BOS 标记的得分。
  • reduce_loss (bool可选) — 仅在传入 labels 时相关。如果为 True,则使用 torch.Tensor.sum 操作对 NLL 损失进行缩减。
  • labels (shape为 (batch_size, sequence_length)torch.LongTensor可选) — 用于计算掩码语言模型损失的标签。索引应在 [0, ..., config.vocab_size] 范围内或为 -100(请参阅 input_ids 文档字符串)。索引设为 -100 的标记将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size] 范围内的标记进行计算。
  • n_docs (int可选) — 要检索的文档数量。

返回

RetrievAugLMMarginOutputtuple(torch.FloatTensor)

一个 RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组(如果传入 return_dict=Falseconfig.return_dict=False 时),根据配置(RagConfig)和输入包含不同的元素。

RagSequenceForGeneration 的前向(forward)方法,重写了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用 Module 实例而不是这个,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。

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

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

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

  • past_key_values (Cache,*可选*,当传入 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。欲了解更多细节,请参阅我们的 KV 缓存指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

示例

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

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

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

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

生成

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.LongTensor | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | None = None doc_scores: torch.FloatTensor | None = None do_deduplication: bool | None = None num_return_sequences: int | None = None num_beams: int | None = None n_docs: int | None = None **model_kwargs ) shape为 (batch_size * num_return_sequences, sequence_length)torch.LongTensor

参数

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

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

    什么是注意力掩码?

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

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

  • doc_scores (shape为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(请参阅 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的得分。

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

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

返回

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

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

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

RagTokenForGeneration

class transformers.RagTokenForGeneration

< >

( config: transformers.configuration_utils.PreTrainedConfig | None = None question_encoder: transformers.modeling_utils.PreTrainedModel | None = None generator: transformers.modeling_utils.PreTrainedModel | None = None retriever: transformers.models.rag.retrieval_rag.RagRetriever | None = None **kwargs )

参数

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

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

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

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.FloatTensor | None = None encoder_outputs: tuple[tuple[torch.Tensor]] | None = None decoder_input_ids: torch.LongTensor | None = None decoder_attention_mask: torch.BoolTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None context_input_ids: torch.LongTensor | None = None context_attention_mask: torch.LongTensor | None = None doc_scores: torch.FloatTensor | None = None use_cache: bool | None = None output_attentions: bool | None = None output_hidden_states: bool | None = None output_retrieved: bool | None = None do_marginalize: bool | None = None reduce_loss: bool | None = None labels: torch.LongTensor | None = None n_docs: int | None = None **kwargs ) RetrievAugLMMarginOutputtuple(torch.FloatTensor)

参数

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

    什么是输入 ID?

  • attention_mask (形状为 (batch_size, sequence_length)torch.FloatTensor, 可选) — 用于避免对填充 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_input_ids (形状为 (batch_size, target_sequence_length)torch.LongTensor, 可选) — 解码器输入序列 token 的词汇表索引。

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

    什么是解码器输入 ID?

  • decoder_attention_mask (形状为 (batch_size, target_sequence_length)torch.BoolTensor, 可选) — 默认行为:生成一个忽略 decoder_input_ids 中 pad token 的张量。因果掩码(Causal mask)也将默认使用。
  • past_key_values (~cache_utils.Cache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由模型在解码的前一个阶段返回的 past_key_values 组成,当 use_cache=Trueconfig.use_cache=True 时使用。

    输入仅允许使用 Cache 实例,请参阅我们的 kv 缓存指南。如果不传递 past_key_values,默认将初始化 DynamicCache

    模型将输出与作为输入馈送的格式相同的缓存。

    如果使用了 past_key_values,用户应仅输入形状为 (batch_size, unprocessed_length) 的未处理 input_ids(那些其过往键值状态未提供给此模型的),而不是形状为 (batch_size, sequence_length) 的所有 input_ids

  • 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__() 返回。
  • doc_scores (形状为 (batch_size, config.n_docs)torch.FloatTensor) — 每个检索到的文档嵌入(参见 retrieved_doc_embeds)与 question_encoder_last_hidden_state 之间的分数。如果模型没有使用 retriever 初始化,则必须将 doc_scores 提供给前向传递。doc_scores 可以通过 question_encoder_last_hidden_stateretrieved_doc_embeds 计算,有关更多信息,请参阅示例。
  • use_cache (bool, 可选) — 如果设置为 True,则返回 past_key_values 键值状态,并可用于加速解码(参见 past_key_values)。
  • output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息,请参阅返回张量下的 attentions
  • output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量下的 hidden_states
  • output_retrieved (bool, 可选) — 是否返回 retrieved_doc_embedsretrieved_doc_idscontext_input_idscontext_attention_mask。有关更多详细信息,请参阅返回张量。
  • do_marginalize (bool, 可选) — 如果为 True,则通过使用 torch.nn.functional.log_softmax 在所有文档上对 logits 进行边际化。
  • reduce_loss (bool, 可选) — 仅在传递了 labels 时相关。如果为 True,则使用 torch.Tensor.sum 操作对 NLL 损失进行缩减。
  • labels (torch.LongTensor,形状为 (batch_size, sequence_length)可选) — 用于计算掩码语言建模损失的标签。索引应在 [0, ..., config.vocab_size] 范围内或为 -100(请参阅 input_ids 文档字符串)。索引设置为 -100 的 token 将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size] 范围内的 token 进行计算。
  • n_docs (int可选) — 要检索的文档数量。

返回

RetrievAugLMMarginOutputtuple(torch.FloatTensor)

一个 RetrievAugLMMarginOutput 或一个 torch.FloatTensor 元组(如果传入 return_dict=Falseconfig.return_dict=False 时),根据配置(RagConfig)和输入包含不同的元素。

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

虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用 Module 实例而不是这个,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。

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

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

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

  • past_key_values (Cache,*可选*,当传入 use_cache=Trueconfig.use_cache=True 时返回) — 这是一个 Cache 实例。欲了解更多细节,请参阅我们的 KV 缓存指南

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

示例

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

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

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

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

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

生成

< >

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

参数

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

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

    什么是注意力掩码?

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

    如果模型未初始化 retriever,则必须在 forward 传递中提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未初始化 retriever,则必须在 forward 传递中提供 context_input_idscontext_input_ids__call__() 返回。

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

    如果模型未初始化 retriever,则必须在 forward 传递中提供 context_input_idscontext_input_ids__call__() 返回。

  • n_docs (int可选,默认为 config.n_docs) — 要检索的文档数量和/或用于生成答案的文档数量。
  • generation_config (~generation.GenerationConfig可选) — 用作生成调用基础参数配置的生成配置。传递给生成并与 generation_config 属性匹配的 **kwargs 将覆盖它们。如果未提供 generation_config,将使用默认配置,其加载优先级如下:1) 如果存在,则从 generation_config.json 模型文件加载;2) 从模型配置中加载。请注意,未指定的参数将继承 GenerationConfig 的默认值,应查阅其文档以了解生成参数化。
  • prefix_allowed_tokens_fn (Callable[[int, torch.Tensor], list[int]]可选) — 如果提供,此函数会将束搜索限制在每一步仅允许的 token 范围内。如果不提供,则不应用约束。此函数接受 2 个参数:inputs_ids 和批处理 ID batch_id。它必须返回一个列表,其中包含根据之前生成的 token inputs_ids 和批处理 ID batch_id 确定的下一个生成步骤的允许 token。此参数对于基于前缀的受限生成非常有用,如 自回归实体检索 中所述。
  • logits_processor (LogitsProcessorList可选) — 自定义 logit 处理器,用于补充根据参数和模型配置构建的默认 logit 处理器。如果传入了已经通过参数或模型配置创建的 logit 处理器,则会抛出错误。
  • 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 解码。

在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.