Transformers documentation
RAG
该模型于 2020-05-22 在 HF papers 上发布,并于 2020-11-16 贡献给 Hugging Face Transformers。
RAG
检索增强生成(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
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0id2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Noneis_encoder_decoder: bool = Truequestion_encoder: transformers.configuration_utils.PreTrainedConfig | dict | None = Nonegenerator: transformers.configuration_utils.PreTrainedConfig | dict | None = Nonevocab_size: int | None = Noneprefix: str | None = Nonebos_token_id: int | None = Nonepad_token_id: int | None = Noneeos_token_id: int | list[int] | None = Nonedecoder_start_token_id: int | None = Nonetitle_sep: str = ' / 'doc_sep: str = ' // 'n_docs: int = 5max_combined_length: int = 300retrieval_vector_size: int = 768retrieval_batch_size: int = 8dataset: str = 'wiki_dpr'dataset_split: str = 'train'index_name: str = 'compressed'index_path: str | None = Nonepassages_path: str | None = Noneuse_dummy_dataset: bool = Falsereduce_loss: bool = Falselabel_smoothing: float = 0.0do_deduplication: bool = Trueexclude_bos_score: bool = Falsedo_marginalize: bool = Falseoutput_retrieved: bool = Falseuse_cache: bool = Truedataset_revision: str | None = None )
参数
- is_encoder_decoder (
bool,可选,默认为True) — 指示该模型是否用作编码器/解码器。 - question_encoder (
PreTrainedConfig | dict,可选) — 问题编码器骨干网络的配置对象或字典 - generator (
PreTrainedConfig | dict,可选) — 答案生成器骨干网络的配置对象或字典 - vocab_size (
int, 可选) — 模型的词汇表大小。定义了input_ids可以表示的不同 token 的数量。 - prefix (
str, 可选) — 在将输入传递给生成器模型之前,添加到每个输入开头的字符串前缀。 - bos_token_id (
int, 可选) — 词汇表中用于表示流起始的 Token id。 - pad_token_id (
int, 可选) — 词表中用于填充的 token id。 - eos_token_id (
Union[int, list[int]], 可选) — 词汇表中用于表示流结束(end-of-stream)的 token id。 - decoder_start_token_id (
int, 可选) — 如果编码器-解码器模型使用不同于bos的 token 开始解码,则指定该 token 的 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,可选,默认为False) — 在计算损失时是否忽略 BOS 标记。 - do_marginalize (
bool,可选,默认为False) — 如果为True,则通过torch.nn.functional.log_softmax对所有文档的 logits 进行边缘化。 - output_retrieved (
bool,可选,默认为False) — 如果设置为True,则会返回retrieved_doc_embeds、retrieved_doc_ids、context_input_ids和context_attention_mask。有关详细信息,请参阅返回的张量。 - use_cache (
bool,可选,默认为True) — 模型是否应返回最后一次键/值注意力(并非所有模型都使用)。仅当config.is_decoder=True或模型为仅解码器生成式模型时相关。 - dataset_revision (
str, 可选,) — 用于检索的 Hugging Face 数据集的版本(commit hash、tag 或 branch)。
这是用于存储 RagModel 配置的类。它根据指定的参数实例化 Rag 模型,并定义模型架构。使用默认值实例化配置将生成与 类似的配置。
配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。
from_question_encoder_generator_configs
< source >( question_encoder_config: PreTrainedConfiggenerator_config: PreTrainedConfig**kwargs ) → EncoderDecoderConfig
配置对象的实例
从预训练的编码器模型配置和 解码器模型配置实例化一个 EncoderDecoderConfig(或其派生类)。
RagTokenizer
Rag 特定输出
class transformers.models.rag.modeling_rag.RetrievAugLMMarginOutput
< source >( loss: typing.Optional[torch.Tensor] = Nonelogits: typing.Optional[torch.Tensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Noneretrieved_doc_embeds: typing.Optional[torch.Tensor] = Noneretrieved_doc_ids: typing.Optional[torch.LongTensor] = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonequestion_encoder_last_hidden_state: typing.Optional[torch.Tensor] = Nonequestion_enc_hidden_states: tuple[torch.Tensor, ...] | None = Nonequestion_enc_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_enc_last_hidden_state: typing.Optional[torch.Tensor] = Nonegenerator_enc_hidden_states: tuple[torch.Tensor, ...] | None = Nonegenerator_enc_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_dec_hidden_states: tuple[torch.Tensor, ...] | None = Nonegenerator_dec_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_cross_attentions: tuple[torch.Tensor, ...] | None = None )
参数
- loss (
torch.Tensor,形状为(1,),可选,在提供labels时返回) — 语言建模损失。 - logits (
torch.Tensor,形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能是针对每个词汇表标记,在所有文档上边缘化后的结果。 - doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。 - past_key_values (
Cache, 可选, 在以下情况下返回use_cache=True被传入或config.use_cache=True) — 它是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv cache 指南.包含解码器中预计算的隐藏状态(注意力块中的键和值),可用于加速顺序解码(参见
past_key_values输入)。 - retrieved_doc_embeds(形状为
(batch_size, config.n_docs, hidden_size)的torch.Tensor,可选,当 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(形状为
(batch_size * config.n_docs, config.max_combined_length)的torch.LongTensor,可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的input_ids进行后处理得到的注意力掩码。 - question_encoder_last_hidden_state (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 问题编码器最后一层输出的隐藏状态序列,即模型的池化输出。 - question_enc_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个对应嵌入层的输出,一个对应每一层的输出),形状为(batch_size, sequence_length, hidden_size).问题编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- question_enc_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).问题编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
- generator_enc_last_hidden_state (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 模型生成器编码器最后一层输出的隐藏状态序列。 - generator_enc_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个对应嵌入层的输出,一个对应每一层的输出),形状为(batch_size, sequence_length, hidden_size).生成器编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- generator_enc_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
- generator_dec_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个用于嵌入的输出,一个用于每一层的输出),形状为(batch_size, sequence_length, hidden_size).生成器解码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- generator_dec_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器解码器的注意力权重,经过注意力 softmax 之后,用于在自注意力头中计算加权平均值。
- generator_cross_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器解码器中交叉注意力的权重,在注意力 softmax 之后,用于在交叉注意力头中计算加权平均值。
检索器增强边际化模型输出的基类。
class transformers.models.rag.modeling_rag.RetrievAugLMOutput
< source >( logits: typing.Optional[torch.Tensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Noneretrieved_doc_embeds: typing.Optional[torch.Tensor] = Noneretrieved_doc_ids: typing.Optional[torch.LongTensor] = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonequestion_encoder_last_hidden_state: typing.Optional[torch.Tensor] = Nonequestion_enc_hidden_states: tuple[torch.Tensor, ...] | None = Nonequestion_enc_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_enc_last_hidden_state: typing.Optional[torch.Tensor] = Nonegenerator_enc_hidden_states: tuple[torch.Tensor, ...] | None = Nonegenerator_enc_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_dec_hidden_states: tuple[torch.Tensor, ...] | None = Nonegenerator_dec_attentions: tuple[torch.Tensor, ...] | None = Nonegenerator_cross_attentions: tuple[torch.Tensor, ...] | None = None )
参数
- logits (
torch.Tensor,形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能是针对每个词汇表标记,在所有文档上边缘化后的结果。 - doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。 - past_key_values (
Cache, 可选, 在以下情况下返回use_cache=True被传入或config.use_cache=True) — 它是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv cache 指南.包含解码器中预计算的隐藏状态(注意力块中的键和值),可用于加速顺序解码(参见
past_key_values输入)。 - retrieved_doc_embeds(形状为
(batch_size, config.n_docs, hidden_size)的torch.Tensor,可选,当 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(形状为
(batch_size * config.n_docs, config.max_combined_length)的torch.LongTensor,可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的input_ids进行后处理得到的注意力掩码。 - question_encoder_last_hidden_state (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 问题编码器最后一层输出的隐藏状态序列,即模型的池化输出。 - question_enc_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个对应嵌入层的输出,一个对应每一层的输出),形状为(batch_size, sequence_length, hidden_size).问题编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- question_enc_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).问题编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
- generator_enc_last_hidden_state (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 模型生成器编码器最后一层输出的隐藏状态序列。 - generator_enc_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个对应嵌入层的输出,一个对应每一层的输出),形状为(batch_size, sequence_length, hidden_size).生成器编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- generator_enc_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
- generator_dec_hidden_states (
tuple(torch.Tensor), 可选, 在传入output_hidden_states=True或config.output_hidden_states=True) — 的元组torch.Tensor((一个用于嵌入的输出,一个用于每一层的输出),形状为(batch_size, sequence_length, hidden_size).生成器解码器在每一层输出处的隐藏状态,以及初始嵌入输出。
- generator_dec_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器解码器的注意力权重,经过注意力 softmax 之后,用于在自注意力头中计算加权平均值。
- generator_cross_attentions (
tuple(torch.Tensor), 可选, 在传入output_attentions=True或config.output_attentions=True) — 的元组torch.Tensor((每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length).生成器解码器中交叉注意力的权重,在注意力 softmax 之后,用于在交叉注意力头中计算加权平均值。
RagRetriever
class transformers.RagRetriever
< source >( configquestion_encoder_tokenizergenerator_tokenizerindex = Noneinit_retrieval = True )
参数
- config (RagConfig) —
此 Retriever 所配合使用的 RAG 模型的配置。包含指示要构建哪个
Index的参数。您可以使用config.index_name="custom"加载自己的自定义数据集,或者例如使用config.index_name="wiki_dpr"从 datasets 库中加载一个标准数据集(默认)。 - question_encoder_tokenizer (PreTrainedTokenizer) — 用于对问题进行分词的 tokenizer。它用于解码问题,然后使用 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")检索器初始化函数。它将索引加载到内存中。
postprocess_docs
< source >( docsinput_stringsprefixn_docsreturn_tensors = None ) → tuple(tensors)
对检索到的 docs 进行后处理,并将其与 input_strings 进行组合。
retrieve
< source >( question_hidden_states: ndarrayn_docs: int ) → tuple[np.ndarray, np.ndarray, list[dict]]
参数
- question_hidden_states(形状为
(batch_size, vector_size)的np.ndarray)— 用于检索的一批查询向量。 - n_docs (
int) — 每个查询检索的文档数量。
返回值
tuple[np.ndarray, np.ndarray, list[dict]]
一个包含以下对象的元组:
- retrieved_doc_embeds(形状为
(batch_size, n_docs, dim)的np.ndarray)— 每个查询所检索文档的检索嵌入。 - doc_ids(形状为
(batch_size, n_docs)的np.ndarray)— 索引中文档的 ID - doc_dicts (
list[dict]):每个查询对应的retrieved_doc_embeds示例。
获取指定 question_hidden_states 的文档。
RagModel
class transformers.RagModel
< source >( config: transformers.configuration_utils.PreTrainedConfig | None = Nonequestion_encoder: transformers.modeling_utils.PreTrainedModel | None = Nonegenerator: transformers.modeling_utils.PreTrainedModel | None = Noneretriever: 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
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneencoder_outputs: tuple[tuple[torch.Tensor]] | None = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.BoolTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Nonedoc_scores: typing.Optional[torch.Tensor] = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Noneuse_cache: bool | None = Noneoutput_attentions: bool | None = Noneoutput_hidden_states: bool | None = Noneoutput_retrieved: bool | None = Nonen_docs: int | None = None**kwargs ) → RetrievAugLMOutput or tuple(torch.Tensor)
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。 RagConfig, 用于初始化模型,指定要使用的生成器,同时指定与之兼容的生成器分词器。请使用该分词器类来获取索引。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token,
- 对于被掩码的 token,值为 0。
- encoder_outputs (
tuple(tuple(torch.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)是生成器编码器最后一层输出的隐藏状态序列。在解码过程中由 (RagModel) 模型使用。
- decoder_input_ids(形状为
(batch_size, target_sequence_length)的torch.LongTensor,可选) — 用于生成任务。默认为None,请按照您使用的 RAG 实例中生成器模型的说明进行构造。 - decoder_input_ids (
torch.LongTensor形状为(batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- decoder_attention_mask(形状为
(batch_size, target_sequence_length)的torch.BoolTensor,可选) — 默认行为:生成一个忽略decoder_input_ids中填充标记的张量。默认情况下也会使用因果掩码。 - past_key_values (
~cache_utils.Cache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由past_key_values在解码的先前阶段返回,当use_cache=True或config.use_cache=True.仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入
past_key_values,则默认会初始化 DynamicCache。模型将输出与输入相同的缓存格式。
如果使用了
past_key_values,用户应仅输入未处理的input_ids(即未向该模型提供其过去键值状态的那些),其形状为(batch_size, unprocessed_length),而不是形状为(batch_size, sequence_length)的全部input_ids。 - doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。如果模型未使用retriever进行初始化,则必须为前向传播提供doc_scores。 可以通过question_encoder_last_hidden_state和retrieved_doc_embeds计算doc_scores,更多信息请参见示例。 - context_input_ids (
torch.LongTensor,形状为(batch_size * config.n_docs, config.max_combined_length),可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的input_ids进行后处理得到的输入 ID。如果模型初始化时未配置retriever,则必须在 forward 过程中提供context_input_ids。context_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_attention_mask。context_attention_mask由__call__()返回。 - use_cache (
bool,可选) — 如果设置为True,则会返回past_key_values键值状态,可用于加速解码(参见past_key_values)。 - output_attentions (
bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息,请参阅返回张量中的attentions。 - output_hidden_states (
bool,可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量中的hidden_states。 - output_retrieved (
bool, 可选) — 是否返回retrieved_doc_embeds、retrieved_doc_ids、context_input_ids和context_attention_mask。有关详细信息,请参阅返回的张量。 - n_docs (
int, 可选) — 要检索的文档数量。
返回值
RetrievAugLMOutput 或 tuple(torch.Tensor)
一个 RetrievAugLMOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(RagConfig)和输入。
RagModel 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义,但之后应调用
Module实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。
logits (
torch.Tensor,形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能是对每个词汇表标记在所有文档上求边际化后得到的。doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。past_key_values(
Cache,可选,在传入use_cache=True或config.use_cache=True时返回)—— 它是一个 Cache 实例。更多详情,请参阅我们的 kv cache 指南。包含解码器中预计算的隐藏状态(注意力块中的键和值),可用于加速顺序解码(参见
past_key_values输入)。retrieved_doc_embeds(形状为
(batch_size, config.n_docs, hidden_size)的torch.Tensor,可选,当 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 (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 问题编码器最后一层输出的隐藏状态序列,即模型的池化输出。question_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每一层的输出),形状为(batch_size, sequence_length, hidden_size)。问题编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
question_enc_attentions (
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回) — 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。问题编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_enc_last_hidden_state(形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor,可选)— 模型生成器编码器最后一层输出的隐藏状态序列。generator_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_enc_attentions(
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_dec_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器解码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_dec_attentions(
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器解码器的注意力权重,经过注意力 softmax 之后,用于在自注意力头中计算加权平均值。
generator_cross_attentions (
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回) — 由torch.Tensor组成的元组(每层一个),形状为(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
< source >( config: transformers.configuration_utils.PreTrainedConfig | None = Nonequestion_encoder: transformers.modeling_utils.PreTrainedModel | None = Nonegenerator: transformers.modeling_utils.PreTrainedModel | None = Noneretriever: 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
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneencoder_outputs: tuple[tuple[torch.Tensor]] | None = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.BoolTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = Noneoutput_attentions: bool | None = Noneoutput_hidden_states: bool | None = Noneoutput_retrieved: bool | None = Noneexclude_bos_score: bool | None = Nonereduce_loss: bool | None = Nonelabels: typing.Optional[torch.LongTensor] = Nonen_docs: int | None = None**kwargs ) → RetrievAugLMMarginOutput or tuple(torch.Tensor)
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。 RagConfig, 用于初始化模型,指定要使用的生成器,同时指定与之兼容的生成器分词器。请使用该分词器类来获取索引。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token,
- 对于被掩码的 token,值为 0。
- encoder_outputs (
tuple(tuple(torch.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)是生成器编码器最后一层输出的隐藏状态序列。在解码过程中由 (RagModel) 模型使用。
- decoder_input_ids(形状为
(batch_size, target_sequence_length)的torch.LongTensor,可选) — 用于生成任务。默认为None,请按照您使用的 RAG 实例中生成器模型的说明进行构造。 - decoder_input_ids (
torch.LongTensor形状为(batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- decoder_attention_mask(形状为
(batch_size, target_sequence_length)的torch.BoolTensor,可选) — 默认行为:生成一个忽略decoder_input_ids中填充标记的张量。默认情况下也会使用因果掩码。 - past_key_values (
~cache_utils.Cache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由past_key_values在解码的先前阶段返回,当use_cache=True或config.use_cache=True.仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入
past_key_values,则默认会初始化 DynamicCache。模型将输出与输入相同的缓存格式。
如果使用了
past_key_values,用户应仅输入未处理的input_ids(即未向该模型提供其过去键值状态的那些),其形状为(batch_size, unprocessed_length),而不是形状为(batch_size, sequence_length)的全部input_ids。 - context_input_ids (
torch.LongTensor,形状为(batch_size * config.n_docs, config.max_combined_length),可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的input_ids进行后处理得到的输入 ID。如果模型初始化时未配置retriever,则必须在 forward 过程中提供context_input_ids。context_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_attention_mask。context_attention_mask由__call__()返回。 - doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。如果模型未使用retriever进行初始化,则必须为前向传播提供doc_scores。 可以通过question_encoder_last_hidden_state和retrieved_doc_embeds计算doc_scores,更多信息请参见示例。 - use_cache (
bool,可选) — 如果设置为True,则会返回past_key_values键值状态,可用于加速解码(参见past_key_values)。 - output_attentions (
bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息,请参阅返回张量中的attentions。 - output_hidden_states (
bool,可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量中的hidden_states。 - output_retrieved (
bool, 可选) — 是否返回retrieved_doc_embeds、retrieved_doc_ids、context_input_ids和context_attention_mask。有关详细信息,请参阅返回的张量。 - exclude_bos_score (
bool, 可选) — 仅在传入labels时相关。如果True,在计算损失时会忽略 BOS 标记的分数。 - reduce_loss (
bool, 可选) — 仅在传入labels时相关。如果为True,则使用torch.Tensor.sum操作对 NLL 损失进行归约。 - labels (
torch.LongTensor,形状为(batch_size, sequence_length),可选) — 用于计算掩码语言建模损失的标签。索引值应在[0, ..., config.vocab_size]范围内,或为 -100(参见input_ids的文档字符串)。索引值设为-100的 token 将被忽略(即被掩码),损失仅针对标签位于[0, ..., config.vocab_size]范围内的 token 进行计算。 - n_docs (
int, 可选) — 要检索的文档数量。
返回值
RetrievAugLMMarginOutput 或 tuple(torch.Tensor)
一个 RetrievAugLMMarginOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(RagConfig)和输入。
RagSequenceForGeneration 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义,但之后应调用
Module实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。
loss(形状为
(1,)的torch.Tensor,可选,在提供labels时返回)— 语言建模损失。logits (
torch.Tensor,形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能是对每个词汇表标记在所有文档上求边际化后得到的。doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。past_key_values(
Cache,可选,在传入use_cache=True或config.use_cache=True时返回)—— 它是一个 Cache 实例。更多详情,请参阅我们的 kv cache 指南。包含解码器中预计算的隐藏状态(注意力块中的键和值),可用于加速顺序解码(参见
past_key_values输入)。retrieved_doc_embeds(形状为
(batch_size, config.n_docs, hidden_size)的torch.Tensor,可选,当 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 (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 问题编码器最后一层输出的隐藏状态序列,即模型的池化输出。question_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每一层的输出),形状为(batch_size, sequence_length, hidden_size)。问题编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
question_enc_attentions (
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回) — 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。问题编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_enc_last_hidden_state(形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor,可选)— 模型生成器编码器最后一层输出的隐藏状态序列。generator_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_enc_attentions(
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_dec_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器解码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_dec_attentions(
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器解码器的注意力权重,经过注意力 softmax 之后,用于在自注意力头中计算加权平均值。
generator_cross_attentions (
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回) — 由torch.Tensor组成的元组(每层一个),形状为(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,
... )generate
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.LongTensor] = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Nonedo_deduplication: bool | None = Nonenum_return_sequences: int | None = Nonenum_beams: int | None = Nonen_docs: int | None = 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 (
torch.Tensor形状为(batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token,
- 对于被掩码的 token,值为 0。
- context_input_ids(形状为
(batch_size * config.n_docs, config.max_combined_length)的torch.LongTensor,可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的 input_ids 进行后处理得到的输入 ID。 - context_attention_mask (
torch.LongTensor形状为(batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True) — 从检索到的文档和问题编码器input_ids中由检索器后处理得到的注意力掩码。如果模型未使用
retriever进行初始化,或者未提供input_ids,则必须为前向传播提供context_input_ids和context_attention_mask。它们由__call__()返回。 - doc_scores (
torch.Tensor形状为(batch_size, config.n_docs)) — 每个检索文档嵌入(参见retrieved_doc_embeds) )与question_encoder_last_hidden_state.如果模型未使用
retriever进行初始化,或者未提供input_ids,则必须为前向传播提供doc_scores。doc_scores由__call__()返回。 - do_deduplication (
bool, 可选) — 是否对给定输入来自不同上下文文档的生成结果进行去重。在使用分布式后端进行训练时,必须将其设置为False。 - num_return_sequences(
int, 可选,默认为 1) — 批次中每个元素独立计算并返回的序列数量。请注意,这不是我们传递给generator的[generate()](/docs/transformers/main/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()。
返回值
形状为 (batch_size * num_return_sequences, sequence_length) 的 torch.LongTensor
生成的序列。第二个维度(序列长度)要么等于 max_length,要么更短,如果所有批次都因 eos_token_id 而提前结束。
实现 RAG 序列的“thorough”解码。有关如何设置其他 generate 输入参数的更多信息,请参阅 generate()` 文档。
RagTokenForGeneration
class transformers.RagTokenForGeneration
< source >( config: transformers.configuration_utils.PreTrainedConfig | None = Nonequestion_encoder: transformers.modeling_utils.PreTrainedModel | None = Nonegenerator: transformers.modeling_utils.PreTrainedModel | None = Noneretriever: transformers.models.rag.retrieval_rag.RagRetriever | None = None**kwargs )
参数
- config (PreTrainedConfig, 可选) — 包含模型所有参数的模型配置类。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。
- question_encoder (
PreTrainedModel,可选) — 负责将问题编码为用于检索的隐藏状态的模型。 - generator (
PreTrainedModel,可选) — 负责基于检索到的文档生成文本的模型。 - retriever (
RagRetriever,可选) — 负责根据编码后的问题从知识库中检索文档的组件。
RAG-token 模型的实现。它在前向传播过程中执行 RAG-token 特有的边缘化操作。
该模型继承自 PreTrainedModel。请查阅父类文档,了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、剪枝头等)。
该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneencoder_outputs: tuple[tuple[torch.Tensor]] | None = Nonedecoder_input_ids: typing.Optional[torch.LongTensor] = Nonedecoder_attention_mask: typing.Optional[torch.BoolTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = Noneoutput_attentions: bool | None = Noneoutput_hidden_states: bool | None = Noneoutput_retrieved: bool | None = Nonedo_marginalize: bool | None = Nonereduce_loss: bool | None = Nonelabels: typing.Optional[torch.LongTensor] = Nonen_docs: int | None = None**kwargs ) → RetrievAugLMMarginOutput or tuple(torch.Tensor)
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 输入序列中各 token 在词表中的索引。 RagConfig, 用于初始化模型,指定要使用的生成器,同时指定与之兼容的生成器分词器。请使用该分词器类来获取索引。 - attention_mask (
torch.Tensor形状为(batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token,
- 对于被掩码的 token,值为 0。
- encoder_outputs (
tuple(tuple(torch.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)是生成器编码器最后一层输出的隐藏状态序列。在解码过程中由 (RagModel) 模型使用。
- decoder_input_ids(形状为
(batch_size, target_sequence_length)的torch.LongTensor,可选) — 用于生成任务。默认为None,请按照您使用的 RAG 实例中生成器模型的说明进行构造。 - decoder_input_ids (
torch.LongTensor形状为(batch_size, target_sequence_length), ,可选) — 解码器输入序列中各 token 在词汇表中的索引。可通过 AutoTokenizer 获取索引。详情请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。什么是解码器输入 ID?可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- decoder_attention_mask(形状为
(batch_size, target_sequence_length)的torch.BoolTensor,可选) — 默认行为:生成一个忽略decoder_input_ids中填充标记的张量。默认情况下也会使用因果掩码。 - past_key_values (
~cache_utils.Cache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由past_key_values在解码的先前阶段返回,当use_cache=True或config.use_cache=True.仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入
past_key_values,则默认会初始化 DynamicCache。模型将输出与输入相同的缓存格式。
如果使用了
past_key_values,用户应仅输入未处理的input_ids(即未向该模型提供其过去键值状态的那些),其形状为(batch_size, unprocessed_length),而不是形状为(batch_size, sequence_length)的全部input_ids。 - context_input_ids (
torch.LongTensor,形状为(batch_size * config.n_docs, config.max_combined_length),可选,当 output_retrieved=True 时返回) — 由检索器对检索到的文档和问题编码器的input_ids进行后处理得到的输入 ID。如果模型初始化时未配置retriever,则必须在 forward 过程中提供context_input_ids。context_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_attention_mask。context_attention_mask由__call__()返回。 - doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。如果模型未使用retriever进行初始化,则必须为前向传播提供doc_scores。 可以通过question_encoder_last_hidden_state和retrieved_doc_embeds计算doc_scores,更多信息请参见示例。 - use_cache (
bool,可选) — 如果设置为True,则会返回past_key_values键值状态,可用于加速解码(参见past_key_values)。 - output_attentions (
bool, 可选) — 是否返回所有注意力层的注意力张量。有关详细信息,请参阅返回张量中的attentions。 - output_hidden_states (
bool,可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量中的hidden_states。 - output_retrieved (
bool, 可选) — 是否返回retrieved_doc_embeds、retrieved_doc_ids、context_input_ids和context_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, 可选) — 要检索的文档数量。
返回值
RetrievAugLMMarginOutput 或 tuple(torch.Tensor)
一个 RetrievAugLMMarginOutput 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(RagConfig)和输入。
RagTokenForGeneration 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义,但之后应调用
Module实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。
loss(形状为
(1,)的torch.Tensor,可选,在提供labels时返回)— 语言建模损失。logits (
torch.Tensor,形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数。该分数可能是对每个词汇表标记在所有文档上求边际化后得到的。doc_scores(形状为
(batch_size, config.n_docs)的torch.Tensor)— 每个检索到的文档嵌入(参见retrieved_doc_embeds)与question_encoder_last_hidden_state之间的分数。past_key_values(
Cache,可选,在传入use_cache=True或config.use_cache=True时返回)—— 它是一个 Cache 实例。更多详情,请参阅我们的 kv cache 指南。包含解码器中预计算的隐藏状态(注意力块中的键和值),可用于加速顺序解码(参见
past_key_values输入)。retrieved_doc_embeds(形状为
(batch_size, config.n_docs, hidden_size)的torch.Tensor,可选,当 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 (
torch.Tensor,形状为(batch_size, sequence_length, hidden_size),可选) — 问题编码器最后一层输出的隐藏状态序列,即模型的池化输出。question_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每一层的输出),形状为(batch_size, sequence_length, hidden_size)。问题编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
question_enc_attentions (
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回) — 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。问题编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_enc_last_hidden_state(形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor,可选)— 模型生成器编码器最后一层输出的隐藏状态序列。generator_enc_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器编码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_enc_attentions(
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器编码器在注意力 softmax 之后的注意力权重,用于在自注意力头中计算加权平均值。
generator_dec_hidden_states (
tuple(torch.Tensor),可选,在传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组(包含嵌入层的输出以及每个层的输出),形状为(batch_size, sequence_length, hidden_size)。生成器解码器在每一层输出处的隐藏状态,以及初始嵌入输出。
generator_dec_attentions(
tuple(torch.Tensor),可选,在传入output_attentions=True或config.output_attentions=True时返回)—— 一个torch.Tensor元组(每层一个),形状为(batch_size, num_heads, sequence_length, sequence_length)。生成器解码器的注意力权重,经过注意力 softmax 之后,用于在自注意力头中计算加权平均值。
generator_cross_attentions (
tuple(torch.Tensor),可选,当传入output_attentions=True或config.output_attentions=True时返回) — 由torch.Tensor组成的元组(每层一个),形状为(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)generate
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.LongTensor] = Nonecontext_input_ids: typing.Optional[torch.LongTensor] = Nonecontext_attention_mask: typing.Optional[torch.LongTensor] = Nonedoc_scores: typing.Optional[torch.Tensor] = Nonen_docs: int | None = Nonegeneration_config: transformers.generation.configuration_utils.GenerationConfig | None = Noneprefix_allowed_tokens_fn: collections.abc.Callable[[int, torch.Tensor], list[int]] | None = Nonelogits_processor: transformers.generation.logits_process.LogitsProcessorList | None = []stopping_criteria: transformers.generation.stopping_criteria.StoppingCriteriaList | None = []**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 (
torch.Tensor形状为(batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自[0, 1]:- 1 表示未被屏蔽的 token,
- 对于被掩码的 token,值为 0。
- context_input_ids (
torch.LongTensor形状为(batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True) — 时返回 — 经过检索器对检索到的文档和input_ids问题编码器的如果模型未使用
retriever进行初始化,则必须为前向传播提供context_input_ids。context_input_ids由__call__()返回。 - context_attention_mask (
torch.LongTensor形状为(batch_size * config.n_docs, config.max_combined_length), 可选, 当 output_retrieved=True) — 从检索到的文档和问题编码器input_ids中经过检索器后处理的注意力掩码。如果模型未使用
retriever进行初始化,则必须为前向传播提供context_input_ids。context_input_ids由__call__()返回。 - doc_scores (
torch.Tensor形状为(batch_size, config.n_docs)) — 每个检索文档嵌入(参见retrieved_doc_embeds) )与question_encoder_last_hidden_state.如果模型未使用
retriever进行初始化,则必须为前向传播提供context_input_ids。context_input_ids由__call__()返回。 - n_docs (
int,可选,默认为config.n_docs) — 要检索的文档数量,和/或要生成答案的文档数量。 - generation_config (
~generation.GenerationConfig, 可选) — 用作生成调用基础参数化的生成配置。传递给 generate 且与generation_config属性匹配的**kwargs将覆盖这些属性。如果未提供generation_config,则使用默认值,其加载优先级如下:1) 如果存在,则从generation_config.json模型文件加载;2) 从模型配置加载。请注意,未指定的参数将继承 GenerationConfig 的默认值,应查阅其文档以参数化生成过程。 - prefix_allowed_tokens_fn (
Callable[[int, torch.Tensor], list[int]], 可选) — 如果提供,此函数会在每一步将束搜索限制为仅使用允许的 token。如果未提供,则不施加任何约束。此函数接受两个参数inputs_ids和批次 IDbatch_id。它必须返回一个列表,其中包含下一步生成所允许的 token,该列表基于先前生成的 tokeninputs_ids和批次 IDbatch_id进行条件约束。此参数对于基于前缀的约束生成非常有用,如 Autoregressive Entity Retrieval 中所述。 - logits_processor (
LogitsProcessorList, 可选) — 自定义的 logits 处理器,用于补充由参数和模型配置构建的默认 logits 处理器。如果传入的 logit 处理器已经由参数或模型配置创建,则会抛出错误。 - stopping_criteria (
StoppingCriteriaList,可选) — 自定义停止条件,用于补充由参数和模型配置构建的默认停止条件。如果传入的停止条件已经由参数或模型配置创建,则会抛出错误。 - kwargs (
dict[str, Any], 可选) —generate_config的临时参数设置,和/或将被转发到模型forward函数的额外模型特定 kwargs。
返回值
形状为 (batch_size * num_return_sequences, sequence_length) 的 torch.LongTensor
生成的序列。第二个维度(sequence_length)要么等于 max_length,要么更短,如果所有批次都因 eos_token_id 而提前结束。
实现 RAG 令牌解码。