Transformers documentation
RAG
このモデルは2020年5月22日にHF papersで公開され、2020年11月16日にHugging Face Transformersに貢献されました。
RAG
検索拡張生成（RAG）は、事前学習済みニューラルリトリーバーを用いて、事前学習済み言語モデル（パラメトリックメモリ）と外部データソース（非パラメトリックメモリ）へのアクセスを組み合わせる手法です。RAGは推論時に関連するパスージを取得し、それらを条件として生成を行います。これにより、回答の事実性が向上し、モデル全体の再学習を行うことなくインデックスを変更するだけで知識を更新できることが多くあります。
すべての元のRAGチェックポイントは、AI at Meta組織の下で確認できます。
このモデルは 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])Quantizationは、重みを低精度で保存することでメモリを削減します。対応するバックエンドについては、Quantizationの概要を参照してください。 以下の例では、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
クラス 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で表現できるトークンの種類数を定義します。 - prefix (
str、任意) — ジェネレーターモデルに渡す前に、すべての入力に先頭へ付加される文字列プレフィックスです。 - bos_token_id (
int, 省略可) — 語彙におけるストリーム開始に使用されるトークン ID。 - pad_token_id (
int, 任意) — 語彙内でパディングに使用されるトークン ID。 - eos_token_id (
Union[int, list[int]], 省略可) — 語彙内でストリーム終了に使用されるトークン 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のどの分割を読み込むかを指定します。 - index_name (
str、任意、デフォルトは"compressed") —datasetに関連付けられたインデックスのインデックス名です。"legacy"、"exact"、"compressed"のいずれかを選択できます。 - index_path (
str, 省略可) — ディスク上のシリアライズされた faiss インデックスへのパス。 - passages_path (
str, 任意) — faiss インデックスと互換性のあるテキストパッセージへのパス。LegacyIndexを使用する場合は必須です。 - use_dummy_dataset (
bool、省略可、デフォルトはFalse) —datasetで指定されたデータセットの「ダミー」バージョンを読み込むかどうか。 - 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 データセットのリビジョン（コミットハッシュ、タグ、またはブランチ）。
これは 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 (
torch.Tensor、shape(batch_size, config.n_docs)) — 取得されたドキュメントの埋め込み（retrieved_doc_embedsを参照）とquestion_encoder_last_hidden_stateの間のスコア。 - past_key_values (
Cache, 省略可能, が返されるのはuse_cache=Trueが渡された場合、またはconfig.use_cache=True) — これは Cache のインスタンスです。詳細については、 KVキャッシュのガイド.デコーダーの事前計算済み hidden-states(アテンションブロックのキーとバリュー)が含まれており、シーケンシャルデコーディングを高速化するために使用できます(
past_key_values入力参照)。 - retrieved_doc_embeds (
torch.Tensor、形状(batch_size, config.n_docs, hidden_size)、任意、output_retrieved=True の場合に返されます)— retriever によって取得された埋め込みドキュメント。question_encoder_last_hidden_stateと組み合わせてdoc_scoresを計算するために使用されます。 - retrieved_doc_ids (
torch.LongTensor、shape(batch_size, config.n_docs)、optional、output_retrieved=True の場合のみ返されます)— retriever によって取得された埋め込みドキュメントのインデックスです。 - context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、optional、output_retrieved=True の場合のみ返されます)— retriever によって、検索されたドキュメントと question encoder の 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.Tensor、形状(batch_size, sequence_length, hidden_size)、任意) — モデルの question encoder の最終層の出力における、プーリングされた出力の隠れ状態のシーケンスです。 - question_enc_hidden_states (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_hidden_states=Trueが渡された場合、またはconfig.output_hidden_states=True) — のタプルtorch.Tensor(（埋め込みの出力用と各層の出力用が1つずつ）で、形状は(batch_size, sequence_length, hidden_size).クエリエンコーダーの各層の出力および初期埋め込み出力における隠れ状態。
- question_enc_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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(（埋め込みの出力と各層の出力それぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size).ジェネレーターのエンコーダーの各層の出力における隠れ状態、および初期埋め込みの出力。
- generator_enc_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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(（埋め込みの出力と各層の出力それぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size).ジェネレーターのデコーダーにおける各層の出力の隠れ状態、および初期埋め込みの出力。
- generator_dec_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(batch_size, num_heads, sequence_length, sequence_length).自己アテンションヘッドにおける加重平均を計算するために使用される、ジェネレーターのデコーダーのアテンション重み（アテンションのsoftmax適用後）。
- generator_cross_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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 (
torch.Tensor、shape(batch_size, config.n_docs)) — 取得されたドキュメントの埋め込み（retrieved_doc_embedsを参照）とquestion_encoder_last_hidden_stateの間のスコア。 - past_key_values (
Cache, 省略可能, が返されるのはuse_cache=Trueが渡された場合、またはconfig.use_cache=True) — これは Cache のインスタンスです。詳細については、 KVキャッシュのガイド.デコーダーの事前計算済み hidden-states(アテンションブロックのキーとバリュー)が含まれており、シーケンシャルデコーディングを高速化するために使用できます(
past_key_values入力参照)。 - retrieved_doc_embeds (
torch.Tensor、形状(batch_size, config.n_docs, hidden_size)、任意、output_retrieved=True の場合に返されます)— retriever によって取得された埋め込みドキュメント。question_encoder_last_hidden_stateと組み合わせてdoc_scoresを計算するために使用されます。 - retrieved_doc_ids (
torch.LongTensor、shape(batch_size, config.n_docs)、optional、output_retrieved=True の場合のみ返されます)— retriever によって取得された埋め込みドキュメントのインデックスです。 - context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、optional、output_retrieved=True の場合のみ返されます)— retriever によって、検索されたドキュメントと question encoder の 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.Tensor、形状(batch_size, sequence_length, hidden_size)、任意) — モデルの question encoder の最終層の出力における、プーリングされた出力の隠れ状態のシーケンスです。 - question_enc_hidden_states (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_hidden_states=Trueが渡された場合、またはconfig.output_hidden_states=True) — のタプルtorch.Tensor(（埋め込みの出力用と各層の出力用が1つずつ）で、形状は(batch_size, sequence_length, hidden_size).クエリエンコーダーの各層の出力および初期埋め込み出力における隠れ状態。
- question_enc_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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(（埋め込みの出力と各層の出力それぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size).ジェネレーターのエンコーダーの各層の出力における隠れ状態、および初期埋め込みの出力。
- generator_enc_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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(（埋め込みの出力と各層の出力それぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size).ジェネレーターのデコーダーにおける各層の出力の隠れ状態、および初期埋め込みの出力。
- generator_dec_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(batch_size, num_heads, sequence_length, sequence_length).自己アテンションヘッドにおける加重平均を計算するために使用される、ジェネレーターのデコーダーのアテンション重み（アテンションのsoftmax適用後）。
- generator_cross_attentions (
tuple(torch.Tensor), 省略可能, が返されるのはoutput_attentions=Trueが渡された場合、またはconfig.output_attentions=True) — のタプルtorch.Tensor(（各レイヤーごとに1つ）で、形状は(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) — 質問のトークン化に使用されたトークナイザーです。質問をデコードし、その後 generator_tokenizer を使用するために用いられます。
- generator_tokenizer (PreTrainedTokenizer) — RagModel のジェネレーター部分で使用されるトークナイザーです。
- index (
Index、省略可能、デフォルトは設定で定義されたもの)— 指定された場合、設定を使用して構築されたものではなく、このインデックスを使用します。
Retrieverは、ベクトリークエリからドキュメントを取得するために使用されます。ドキュメントの埋め込みとドキュメントの内容を取得し、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")Retrieverの初期化関数です。インデックスをメモリに読み込みます。
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 (
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、shape(batch_size, n_docs)) — インデックス内のドキュメントの 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、任意) — エンコードされた質問に基づいてナレッジベースからドキュメントを取得するコンポーネントです。
特定のヘッドを上に載せずに、生のhidden-statesを出力する素のRag Model。
このモデルは 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)) — 語彙における入力シーケンスのトークンインデックス。 RagConfig, モデルの初期化に使用され、使用するジェネレーターを指定します。また、互換性のあるジェネレーター用トークナイザーも指定します。そのトークナイザークラスを使用してインデックスを取得します。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 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)は、generator のエンコーダーの最終層の出力における隠れ状態のシーケンスです。デコーディング中に(RagModel)モデルで使用されます。
- decoder_input_ids (
torch.LongTensor、shape(batch_size, target_sequence_length)、任意) — 生成タスクで指定します。デフォルトはNoneです。RAG インスタンスで使用している生成モデルの指示に従って構築してください。 - decoder_input_ids (
torch.LongTensorの形状(batch_size, target_sequence_length), 、任意) — 語彙におけるデコーダー入力シーケンスのトークンインデックスです。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- decoder_attention_mask (
torch.BoolTensor、shape(batch_size, target_sequence_length)、任意) — デフォルトの動作:decoder_input_ids内の pad トークンを無視するテンソルを生成します。デフォルトでは、causal mask も使用されます。 - 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, sequence_length)）ではなく、このモデルに過去のキー・バリュー状態が与えられていない未処理のinput_ids（形状(batch_size, unprocessed_length)）のみを入力することになります。 - doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) —retrieved_doc_embeds（取得されたドキュメントの埋め込み、詳細はretrieved_doc_embedsを参照）とquestion_encoder_last_hidden_stateの間のスコアです。モデルがretrieverを使用して初期化されていない場合、doc_scoresを forward pass に提供する必要があります。doc_scoresはquestion_encoder_last_hidden_stateとretrieved_doc_embedsを用いて計算できます。詳細は例を参照してください。 - context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと question encoder のinput_idsから後処理された入力 ID です。モデルがretrieverを用いて初期化されていない場合、forward pass にcontext_input_idsを提供する必要があります。context_input_idsは__call__()によって返されます。 - context_attention_mask (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと質問エンコーダーのinput_idsから後処理されたアテンションマスクです。モデルがretrieverで初期化されていない場合、forward pass に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 を返すかどうかを指定します。詳細は、返されるテンソルのhidden_statesを参照してください。 - output_retrieved (
bool、省略可) —retrieved_doc_embeds、retrieved_doc_ids、context_input_idsおよびcontext_attention_maskを返すかどうかを指定します。詳細は返されるテンソルを参照してください。 - n_docs (
int, 省略可) — 取得するドキュメントの数。
戻り値
RetrievAugLMOutput または tuple(torch.Tensor)
torch.Tensor のタプル、または RetrievAugLMOutput（return_dict=False が渡された場合、または config.return_dict=False の場合）。設定（RagConfig）や入力に応じて、さまざまな要素で構成されます。
RagModelのforwardメソッドは、__call__特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
logits (
torch.Tensor、形状(batch_size, sequence_length, config.vocab_size)) — 言語モデリングヘッドの予測スコアです。各語彙トークンについて、スコアはすべての文書にわたってマージナライズされている可能性があります。doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) — 各検索ドキュメントの埋め込み(retrieved_doc_embedsを参照)とquestion_encoder_last_hidden_stateの間のスコア。past_key_values (
Cache、任意、use_cache=Trueが渡された場合またはconfig.use_cache=Trueの場合に返されます) — Cache のインスタンスです。詳細については、kv cache ガイドを参照してください。デコーダーの事前計算済み hidden-states(アテンションブロックのキーとバリュー)が含まれており、シーケンシャルデコーディングを高速化するために使用できます(
past_key_values入力参照)。retrieved_doc_embeds (
torch.Tensor、形状(batch_size, config.n_docs, hidden_size)、任意、output_retrieved=True の場合に返されます) — リトリーバーによって取得された埋め込みドキュメント。question_encoder_last_hidden_stateと組み合わせてdoc_scoresを計算するために使用されます。retrieved_doc_ids (
torch.LongTensor、shape(batch_size, config.n_docs)、optional、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、shape(batch_size * config.n_docs, config.max_combined_length)、optional、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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力と各層の出力それぞれに対応するもの）で、形状は(batch_size, sequence_length, hidden_size)です。クエリエンコーダーの各層の出力および初期埋め込み出力における隠れ状態。
question_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます) —torch.Tensorのタプル(埋め込みの出力と各レイヤーの出力それぞれに対応するもの)で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのエンコーダーの各層の出力における隠れ状態、および初期埋め込みの出力。
generator_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(batch_size, num_heads, sequence_length, sequence_length)。アテンションのsoftmaxを適用した後の、ジェネレーターのエンコーダーのアテンション重み。自己アテンションヘッドにおける加重平均の計算に使用されます。
generator_dec_hidden_states (
tuple(torch.Tensor)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力用と各レイヤーの出力用がそれぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのデコーダーにおける各層の出力の隠れ状態、および初期埋め込みの出力。
generator_dec_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル(レイヤーごとに1つ)、形状は(batch_size, num_heads, sequence_length, sequence_length)。自己アテンションヘッドにおける加重平均を計算するために使用される、ジェネレーターのデコーダーのアテンション重み（アテンションのsoftmax適用後）。
generator_cross_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます)—torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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モデルの実装です。forward passにおいて、RAG-sequence固有の周辺化（marginalization）を実行します。
このモデルは 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)) — 語彙における入力シーケンスのトークンインデックス。 RagConfig, モデルの初期化に使用され、使用するジェネレーターを指定します。また、互換性のあるジェネレーター用トークナイザーも指定します。そのトークナイザークラスを使用してインデックスを取得します。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 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)は、generator のエンコーダーの最終層の出力における隠れ状態のシーケンスです。デコーディング中に(RagModel)モデルで使用されます。
- decoder_input_ids (
torch.LongTensor、shape(batch_size, target_sequence_length)、任意) — 生成タスクで指定します。デフォルトはNoneです。RAG インスタンスで使用している生成モデルの指示に従って構築してください。 - decoder_input_ids (
torch.LongTensorの形状(batch_size, target_sequence_length), 、任意) — 語彙におけるデコーダー入力シーケンスのトークンインデックスです。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- decoder_attention_mask (
torch.BoolTensor、shape(batch_size, target_sequence_length)、任意) — デフォルトの動作:decoder_input_ids内の pad トークンを無視するテンソルを生成します。デフォルトでは、causal mask も使用されます。 - 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, sequence_length)）ではなく、このモデルに過去のキー・バリュー状態が与えられていない未処理のinput_ids（形状(batch_size, unprocessed_length)）のみを入力することになります。 - context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと question encoder のinput_idsから後処理された入力 ID です。モデルがretrieverを用いて初期化されていない場合、forward pass にcontext_input_idsを提供する必要があります。context_input_idsは__call__()によって返されます。 - context_attention_mask (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと質問エンコーダーのinput_idsから後処理されたアテンションマスクです。モデルがretrieverで初期化されていない場合、forward pass にcontext_attention_maskを提供する必要があります。context_attention_maskは__call__()によって返されます。 - doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) —retrieved_doc_embeds（取得されたドキュメントの埋め込み、詳細はretrieved_doc_embedsを参照）とquestion_encoder_last_hidden_stateの間のスコアです。モデルがretrieverを使用して初期化されていない場合、doc_scoresを forward pass に提供する必要があります。doc_scoresはquestion_encoder_last_hidden_stateとretrieved_doc_embedsを用いて計算できます。詳細は例を参照してください。 - use_cache (
bool、任意) —Trueに設定すると、past_key_valuesのキーバリュー状態が返され、デコーディングの高速化に利用できます（past_key_valuesを参照）。 - output_attentions (
bool, 省略可) — すべてのアテンション層のアテンションテンソルを返すかどうかを指定します。詳細は、返されるテンソルのattentionsを参照してください。 - output_hidden_states (
bool, 省略可) — すべての層の hidden states を返すかどうかを指定します。詳細は、返されるテンソルの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の場合、NLL ロスはtorch.Tensor.sum操作によって集約されます。 - labels (
torch.LongTensor、shape(batch_size, sequence_length)、任意) — マスク言語モデリング損失を計算するためのラベルです。インデックスは[0, ..., config.vocab_size]の範囲内であるか、-100 である必要があります（input_idsの docstring を参照）。インデックスが-100に設定されたトークンは無視（マスク）され、損失は[0, ..., config.vocab_size]の範囲内のラベルを持つトークンに対してのみ計算されます。 - n_docs (
int, 省略可) — 取得するドキュメントの数。
戻り値
RetrievAugLMMarginOutput または tuple(torch.Tensor)
torch.Tensor のタプル（return_dict=False が渡された場合、または config.return_dict=False の場合）または RetrievAugLMMarginOutput。設定（RagConfig）と入力に応じて、さまざまな要素で構成されます。
RagSequenceForGenerationのforwardメソッドは、__call__特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
loss (
torch.Tensor、形状(1,)、任意、labelsが指定された場合に返されます) — 言語モデリングの損失。logits (
torch.Tensor、形状(batch_size, sequence_length, config.vocab_size)) — 言語モデリングヘッドの予測スコアです。各語彙トークンについて、スコアはすべての文書にわたってマージナライズされている可能性があります。doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) — 各検索ドキュメントの埋め込み(retrieved_doc_embedsを参照)とquestion_encoder_last_hidden_stateの間のスコア。past_key_values (
Cache、任意、use_cache=Trueが渡された場合またはconfig.use_cache=Trueの場合に返されます) — Cache のインスタンスです。詳細については、kv cache ガイドを参照してください。デコーダーの事前計算済み hidden-states(アテンションブロックのキーとバリュー)が含まれており、シーケンシャルデコーディングを高速化するために使用できます(
past_key_values入力参照)。retrieved_doc_embeds (
torch.Tensor、形状(batch_size, config.n_docs, hidden_size)、任意、output_retrieved=True の場合に返されます) — リトリーバーによって取得された埋め込みドキュメント。question_encoder_last_hidden_stateと組み合わせてdoc_scoresを計算するために使用されます。retrieved_doc_ids (
torch.LongTensor、shape(batch_size, config.n_docs)、optional、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、shape(batch_size * config.n_docs, config.max_combined_length)、optional、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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力と各層の出力それぞれに対応するもの）で、形状は(batch_size, sequence_length, hidden_size)です。クエリエンコーダーの各層の出力および初期埋め込み出力における隠れ状態。
question_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます) —torch.Tensorのタプル(埋め込みの出力と各レイヤーの出力それぞれに対応するもの)で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのエンコーダーの各層の出力における隠れ状態、および初期埋め込みの出力。
generator_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(batch_size, num_heads, sequence_length, sequence_length)。アテンションのsoftmaxを適用した後の、ジェネレーターのエンコーダーのアテンション重み。自己アテンションヘッドにおける加重平均の計算に使用されます。
generator_dec_hidden_states (
tuple(torch.Tensor)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力用と各レイヤーの出力用がそれぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのデコーダーにおける各層の出力の隠れ状態、および初期埋め込みの出力。
generator_dec_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル(レイヤーごとに1つ)、形状は(batch_size, num_heads, sequence_length, sequence_length)。自己アテンションヘッドにおける加重平均を計算するために使用される、ジェネレーターのデコーダーのアテンション重み（アテンションのsoftmax適用後）。
generator_cross_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます)—torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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 (
torch.LongTensor、shape(batch_size, sequence_length)、任意) — 生成のプロンプトとして使用されるシーケンスです。input_idsが渡されない場合、context_input_idsを指定する必要があります。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
- context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、optional、output_retrieved=True の場合に返される)— retriever によって、検索されたドキュメントと質問エンコーダーの input_ids から後処理された入力 ID。 - context_attention_mask (
torch.LongTensorの形状(batch_size * config.n_docs, config.max_combined_length), 、省略可能, が返されるのは の場合) — 検索されたドキュメントと質問エンコーダーから後処理されたアテンションマスク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が指定されていない場合は、forward pass にdoc_scoresを指定する必要があります。doc_scoresは__call__()によって返されます。 - do_deduplication (
bool, 省略可) — 与えられた入力に対して、異なるコンテキスト文書からの生成結果を重複排除するかどうかを指定します。分散バックエンドでトレーニング中に使用する場合は、Falseに設定する必要があります。 - num_return_sequences(
int, optional, デフォルト値は 1) — バッチ内の各要素に対して、独立して計算して返されるシーケンスの数です。なお、これはgeneratorの[generate()](/docs/transformers/v5.23.0/en/main_classes/text_generation#transformers.GenerationMixin.generate)関数に渡す値ではなく、ここではnum_return_sequencesをnum_beamsに設定しています。 - num_beams (
int、任意、デフォルトは 1) — ビームサーチに使用するビーム数です。1 を指定すると、ビームサーチを行いません。 - n_docs (
int、任意、デフォルトはconfig.n_docs) — 取得するドキュメントの数、および/または回答を生成する対象のドキュメントの数。 - kwargs (
dict[str, Any]、省略可) — 追加の kwargs は generate() に渡されます。
戻り値
torch.LongTensor、形状 (batch_size * num_return_sequences, sequence_length)
生成されたシーケンス。2 次元目（シーケンス長）は、max_length と等しいか、すべてのバッチがeos_token_idにより早期に終了した場合はそれより短くなります。
RAGシーケンスの「thorough」デコーディングを実装します。他のgenerate入力パラメータの設定方法については、generate()`のドキュメントを参照してください。
RagTokenForGeneration
クラス 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トークンモデルの実装です。フォワードパスにおいて、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 = 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)) — 語彙における入力シーケンスのトークンインデックス。 RagConfig, モデルの初期化に使用され、使用するジェネレーターを指定します。また、互換性のあるジェネレーター用トークナイザーも指定します。そのトークナイザークラスを使用してインデックスを取得します。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 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)は、generator のエンコーダーの最終層の出力における隠れ状態のシーケンスです。デコーディング中に(RagModel)モデルで使用されます。
- decoder_input_ids (
torch.LongTensor、shape(batch_size, target_sequence_length)、任意) — 生成タスクで指定します。デフォルトはNoneです。RAG インスタンスで使用している生成モデルの指示に従って構築してください。 - decoder_input_ids (
torch.LongTensorの形状(batch_size, target_sequence_length), 、任意) — 語彙におけるデコーダー入力シーケンスのトークンインデックスです。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- decoder_attention_mask (
torch.BoolTensor、shape(batch_size, target_sequence_length)、任意) — デフォルトの動作:decoder_input_ids内の pad トークンを無視するテンソルを生成します。デフォルトでは、causal mask も使用されます。 - 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, sequence_length)）ではなく、このモデルに過去のキー・バリュー状態が与えられていない未処理のinput_ids（形状(batch_size, unprocessed_length)）のみを入力することになります。 - context_input_ids (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと question encoder のinput_idsから後処理された入力 ID です。モデルがretrieverを用いて初期化されていない場合、forward pass にcontext_input_idsを提供する必要があります。context_input_idsは__call__()によって返されます。 - context_attention_mask (
torch.LongTensor、shape(batch_size * config.n_docs, config.max_combined_length)、任意、output_retrieved=True の場合に返されます)— retriever によって、検索されたドキュメントと質問エンコーダーのinput_idsから後処理されたアテンションマスクです。モデルがretrieverで初期化されていない場合、forward pass にcontext_attention_maskを提供する必要があります。context_attention_maskは__call__()によって返されます。 - doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) —retrieved_doc_embeds（取得されたドキュメントの埋め込み、詳細はretrieved_doc_embedsを参照）とquestion_encoder_last_hidden_stateの間のスコアです。モデルがretrieverを使用して初期化されていない場合、doc_scoresを forward pass に提供する必要があります。doc_scoresはquestion_encoder_last_hidden_stateとretrieved_doc_embedsを用いて計算できます。詳細は例を参照してください。 - use_cache (
bool、任意) —Trueに設定すると、past_key_valuesのキーバリュー状態が返され、デコーディングの高速化に利用できます（past_key_valuesを参照）。 - output_attentions (
bool, 省略可) — すべてのアテンション層のアテンションテンソルを返すかどうかを指定します。詳細は、返されるテンソルのattentionsを参照してください。 - output_hidden_states (
bool, 省略可) — すべての層の hidden states を返すかどうかを指定します。詳細は、返されるテンソルの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の場合、NLL ロスはtorch.Tensor.sum操作によって集約されます。 - labels (
torch.LongTensor、shape(batch_size, sequence_length)、任意) — マスク言語モデリング損失を計算するためのラベルです。インデックスは[0, ..., config.vocab_size]の範囲内であるか、-100 である必要があります（input_idsの docstring を参照）。インデックスが-100に設定されたトークンは無視（マスク）され、損失は[0, ..., config.vocab_size]の範囲内のラベルを持つトークンに対してのみ計算されます。 - n_docs (
int, 省略可) — 取得するドキュメントの数。
戻り値
RetrievAugLMMarginOutput または tuple(torch.Tensor)
torch.Tensor のタプル（return_dict=False が渡された場合、または config.return_dict=False の場合）または RetrievAugLMMarginOutput。設定（RagConfig）と入力に応じて、さまざまな要素で構成されます。
RagTokenForGenerationのforwardメソッドは、__call__特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
loss (
torch.Tensor、形状(1,)、任意、labelsが指定された場合に返されます) — 言語モデリングの損失。logits (
torch.Tensor、形状(batch_size, sequence_length, config.vocab_size)) — 言語モデリングヘッドの予測スコアです。各語彙トークンについて、スコアはすべての文書にわたってマージナライズされている可能性があります。doc_scores (
torch.Tensor、shape(batch_size, config.n_docs)) — 各検索ドキュメントの埋め込み(retrieved_doc_embedsを参照)とquestion_encoder_last_hidden_stateの間のスコア。past_key_values (
Cache、任意、use_cache=Trueが渡された場合またはconfig.use_cache=Trueの場合に返されます) — Cache のインスタンスです。詳細については、kv cache ガイドを参照してください。デコーダーの事前計算済み hidden-states(アテンションブロックのキーとバリュー)が含まれており、シーケンシャルデコーディングを高速化するために使用できます(
past_key_values入力参照)。retrieved_doc_embeds (
torch.Tensor、形状(batch_size, config.n_docs, hidden_size)、任意、output_retrieved=True の場合に返されます) — リトリーバーによって取得された埋め込みドキュメント。question_encoder_last_hidden_stateと組み合わせてdoc_scoresを計算するために使用されます。retrieved_doc_ids (
torch.LongTensor、shape(batch_size, config.n_docs)、optional、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、shape(batch_size * config.n_docs, config.max_combined_length)、optional、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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力と各層の出力それぞれに対応するもの）で、形状は(batch_size, sequence_length, hidden_size)です。クエリエンコーダーの各層の出力および初期埋め込み出力における隠れ状態。
question_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます) —torch.Tensorのタプル(埋め込みの出力と各レイヤーの出力それぞれに対応するもの)で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのエンコーダーの各層の出力における隠れ状態、および初期埋め込みの出力。
generator_enc_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(batch_size, num_heads, sequence_length, sequence_length)。アテンションのsoftmaxを適用した後の、ジェネレーターのエンコーダーのアテンション重み。自己アテンションヘッドにおける加重平均の計算に使用されます。
generator_dec_hidden_states (
tuple(torch.Tensor)、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます)—torch.Tensorのタプル（埋め込みの出力用と各レイヤーの出力用がそれぞれ1つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。ジェネレーターのデコーダーにおける各層の出力の隠れ状態、および初期埋め込みの出力。
generator_dec_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます) —torch.Tensorのタプル(レイヤーごとに1つ)、形状は(batch_size, num_heads, sequence_length, sequence_length)。自己アテンションヘッドにおける加重平均を計算するために使用される、ジェネレーターのデコーダーのアテンション重み（アテンションのsoftmax適用後）。
generator_cross_attentions (
tuple(torch.Tensor)、optional、output_attentions=Trueが渡された場合またはconfig.output_attentions=Trueの場合に返されます)—torch.Tensorのタプル（各レイヤーごとに1つ）、形状は(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 (
torch.LongTensor、shape(batch_size, sequence_length)、任意) — 生成のプロンプトとして使用されるシーケンスです。input_idsが渡されない場合、context_input_idsを指定する必要があります。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
- context_input_ids (
torch.LongTensorの形状(batch_size * config.n_docs, config.max_combined_length), 、オプション, が返されるのは が True のとき) — 検索されたドキュメントと質問エンコーダーのinput_idsをリトリーバーが後処理した入力 ID。モデルが
retrieverで初期化されていない場合、forward pass にcontext_input_idsを指定する必要があります。context_input_idsは__call__()によって返されます。 - context_attention_mask (
torch.LongTensorの形状(batch_size * config.n_docs, config.max_combined_length), オプション, が返されるのは が True の場合) — 検索されたドキュメントと質問エンコーダーから後処理されたアテンションマスクinput_idsをリトリーバーによって。モデルが
retrieverで初期化されていない場合、forward pass に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で初期化されていない場合、forward pass にcontext_input_idsを指定する必要があります。context_input_idsは__call__()によって返されます。 - n_docs (
int、任意、デフォルトはconfig.n_docs) — 取得するドキュメントの数、および/または回答を生成する対象のドキュメントの数。 - generation_config (
~generation.GenerationConfig, 任意) — 生成呼び出しのベースパラメータ化として使用される生成設定です。**kwargsとして generate に渡され、generation_configの属性と一致するものは、それらをオーバーライドします。generation_configが指定されない場合、デフォルトが使用されます。その読み込み優先順位は以下の通りです: 1)generation_config.jsonモデルファイルが存在する場合、そこから読み込む; 2) モデル設定から読み込む。未指定のパラメータは GenerationConfig のデフォルト値を継承することに注意してください。生成をパラメータ化するには、そのドキュメントを確認する必要があります。 - prefix_allowed_tokens_fn (
Callable[[int, torch.Tensor], list[int]], 省略可) — 指定された場合、この関数はビームサーチを各ステップで許可されたトークンのみに制限します。指定されない場合、制限は適用されません。この関数は 2 つの引数inputs_idsとバッチ IDbatch_idを受け取ります。以前に生成されたトークンinputs_idsとバッチ IDbatch_idに条件付けられた次の生成ステップで許可されるトークンのリストを返す必要があります。この引数は、Autoregressive Entity Retrieval に記載されているような、プレフィックスに条件付けられた制約付き生成に有用です。 - logits_processor (
LogitsProcessorList, 任意) — 引数やモデルの config から構築されるデフォルトの logits processors を補完する、カスタムの logits processors です。引数やモデルの config で既に作成済みの logit processor が渡された場合、エラーが発生します。 - stopping_criteria (
StoppingCriteriaList, 任意) — 引数やモデルの config から構築されるデフォルトの停止条件を補完する、カスタムの停止条件です。引数やモデルの config で既に作成済みの停止条件が渡された場合、エラーが発生します。 - kwargs (
dict[str, Any]、省略可) —generate_configのアドホックなパラメータ設定、および/またはモデルのforward関数に転送される追加のモデル固有の kwargs。
戻り値
torch.LongTensor、形状 (batch_size * num_return_sequences, sequence_length)
生成されたシーケンス。2 次元目（sequence_length）は、eos_token_id によりすべてのバッチが早期に終了した場合を除き、max_length と等しくなります。
RAGトークンデコーディングを実装します。