



"""simple docstring"""






import logging
import os
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import List, Optional

import faiss
import torch
from datasets import Features, Sequence, Value, load_dataset

from transformers import (
    DPRContextEncoder,
    DPRContextEncoderTokenizerFast,
    HfArgumentParser,
    RagRetriever,
    RagSequenceForGeneration,
    RagTokenizer,
)


__UpperCAmelCase      :      Union[str, Any]   					=					logging.getLogger(__name__)
torch.set_grad_enabled(False)
__UpperCAmelCase      :      Optional[Any]   					=					'cuda' if torch.cuda.is_available() else 'cpu'
def      A      (			_A,     _A=100,     _A=" "							):




					"""simple docstring"""
					snake_case_						:Any								=						text.split(_A							)
					return [character.join(text[i : i + n]							).strip() for i in range(0,     len(_A							),     _A							)]
def      A      (			_A							):




					"""simple docstring"""
					snake_case_	,  snake_case_						:Dict								=						[], []
					for title, text in zip(documents["title"],     documents["text"]							):
										if text is not None:
															for passage in split_text(_A							):
																				titles.append(title if title is not None else ""							)
																				texts.append(_A							)
					return {"title": titles, "text": texts}
def      A      (			_A,     _A,     _A							):




					"""simple docstring"""
					snake_case_						:int								=						ctx_tokenizer(
					    documents["title"],     documents["text"],     truncation=_A,     padding="longest",     return_tensors="pt"							)["input_ids"]
					snake_case_						:Tuple								=						ctx_encoder(input_ids.to(device=_A							),     return_dict=_A							).pooler_output
					return {"embeddings": embeddings.detach().cpu().numpy()}
def      A      (			_A,     _A,     _A,     ):




					"""simple docstring"""
					######################################
					logger.info("Step 1 - Create the dataset"							)
					######################################

					# The dataset needed for RAG must have three columns:
					# - title (string): title of the document
					# - text (string): text of a passage of the document
					# - embeddings (array of dimension d): DPR representation of the passage

					# Let's say you have documents in tab-separated csv files with columns "title" and "text"
					assert os.path.isfile(rag_example_args.csv_path							), "Please provide a valid path to a csv file"

					# You can load a Dataset object this way
					snake_case_						:int								=						load_dataset(
					    "csv",     data_files=[rag_example_args.csv_path],     split="train",     delimiter="\t",     column_names=["title", "text"]							)

					# More info about loading csv files in the documentation: https://huggingface.co/docs/datasets/loading_datasets.html?highlight=csv#csv-files

					# Then split the documents into passages of 100 words
					snake_case_						:Optional[int]								=						dataset.map(_A,     batched=_A,     num_proc=processing_args.num_proc							)

					# And compute the embeddings
					snake_case_						:str								=						DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name							).to(device=_A							)
					snake_case_						:Optional[Any]								=						DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name							)
					snake_case_						:List[Any]								=						Features(
					    {"text": Value("string"							), "title": Value("string"							), "embeddings": Sequence(Value("float32"							)							)}							)  # optional, save as float32 instead of float64 to save space
					snake_case_						:Any								=						dataset.map(
					    partial(_A,     ctx_encoder=_A,     ctx_tokenizer=_A							),     batched=_A,     batch_size=processing_args.batch_size,     features=_A,     )

					# And finally save your dataset
					snake_case_						:Dict								=						os.path.join(rag_example_args.output_dir,     "my_knowledge_dataset"							)
					dataset.save_to_disk(_A							)
					# from datasets import load_from_disk
					# dataset = load_from_disk(passages_path)  # to reload the dataset

					######################################
					logger.info("Step 2 - Index the dataset"							)
					######################################

					# Let's use the Faiss implementation of HNSW for fast approximate nearest neighbor search
					snake_case_						:int								=						faiss.IndexHNSWFlat(index_hnsw_args.d,     index_hnsw_args.m,     faiss.METRIC_INNER_PRODUCT							)
					dataset.add_faiss_index("embeddings",     custom_index=_A							)

					# And save the index
					snake_case_						:List[Any]								=						os.path.join(rag_example_args.output_dir,     "my_knowledge_dataset_hnsw_index.faiss"							)
					dataset.get_index("embeddings"							).save(_A							)
					# dataset.load_faiss_index("embeddings", index_path)  # to reload the index

					######################################
					logger.info("Step 3 - Load RAG"							)
					######################################

					# Easy way to load the model
					snake_case_						:Optional[int]								=						RagRetriever.from_pretrained(
					    rag_example_args.rag_model_name,     index_name="custom",     indexed_dataset=_A							)
					snake_case_						:str								=						RagSequenceForGeneration.from_pretrained(rag_example_args.rag_model_name,     retriever=_A							)
					snake_case_						:List[Any]								=						RagTokenizer.from_pretrained(rag_example_args.rag_model_name							)

					# For distributed fine-tuning you'll need to provide the paths instead, as the dataset and the index are loaded separately.
					# retriever = RagRetriever.from_pretrained(rag_model_name, index_name="custom", passages_path=passages_path, index_path=index_path)

					######################################
					logger.info("Step 4 - Have fun"							)
					######################################

					snake_case_						:Tuple								=						rag_example_args.question or "What does Moses' rod turn into ?"
					snake_case_						:Any								=						tokenizer.question_encoder(_A,     return_tensors="pt"							)["input_ids"]
					snake_case_						:Any								=						model.generate(_A							)
					snake_case_						:Dict								=						tokenizer.batch_decode(_A,     skip_special_tokens=_A							)[0]
					logger.info("Q: " + question							)
					logger.info("A: " + generated_string							)

@dataclass
class 	__lowerCAmelCase							:


					'''simple docstring'''


					a__        =     field(
					    default=str(Path(__UpperCamelCase  ).parent / 'test_data' / 'my_knowledge_dataset.csv'  )					,metadata={'help': 'Path to a tab-separated csv file with columns \'title\' and \'text\''}					,)
					a__        =     field(
					    default=__UpperCamelCase					,metadata={'help': 'Question that is passed as input to RAG. Default is \'What does Moses\' rod turn into ?\'.'}					,)
					a__        =     field(
					    default='facebook/rag-sequence-nq'					,metadata={'help': 'The RAG model to use. Either \'facebook/rag-sequence-nq\' or \'facebook/rag-token-nq\''}					,)
					a__        =     field(
					    default='facebook/dpr-ctx_encoder-multiset-base'					,metadata={
					        'help': (
					            'The DPR context encoder model to use. Either \'facebook/dpr-ctx_encoder-single-nq-base\' or'
					            ' \'facebook/dpr-ctx_encoder-multiset-base\''
					        )
					    }					,)
					a__        =     field(
					    default=__UpperCamelCase					,metadata={'help': 'Path to a directory where the dataset passages and the index will be saved'}					,)



@dataclass
class 	__lowerCAmelCase							:


					'''simple docstring'''


					a__        =     field(
					    default=__UpperCamelCase					,metadata={
					        'help': 'The number of processes to use to split the documents into passages. Default is single process.'
					    }					,)
					a__        =     field(
					    default=16					,metadata={
					        'help': 'The batch size to use when computing the passages embeddings using the DPR context encoder.'
					    }					,)



@dataclass
class 	__lowerCAmelCase							:


					'''simple docstring'''


					a__        =     field(
					    default=7_68					,metadata={'help': 'The dimension of the embeddings to pass to the HNSW Faiss index.'}					,)
					a__        =     field(
					    default=1_28					,metadata={
					        'help': (
					            'The number of bi-directional links created for every new element during the HNSW index construction.'
					        )
					    }					,)


if __name__ == "__main__":
				logging.basicConfig(level=logging.WARNING)
				logger.setLevel(logging.INFO)

				__UpperCAmelCase      :      int   					=					HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
				__UpperCAmelCase					,	__UpperCAmelCase					,	__UpperCAmelCase      :      List[str]   					=					parser.parse_args_into_dataclasses()
				with TemporaryDirectory() as tmp_dir:
								__UpperCAmelCase      :      str   					=					rag_example_args.output_dir or tmp_dir
								main(rag_example_args, processing_args, index_hnsw_args)






"""simple docstring"""






from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__UpperCAmelCase      :      Union[str, Any]   					=					logging.get_logger(__name__)

__UpperCAmelCase      :      str   					=					{
    'google/bigbird-pegasus-large-arxiv': (
        'https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-pubmed': (
        'https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-bigpatent': (
        'https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json'
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}

class 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					a__        =     'bigbird_pegasus'
					a__        =     ['past_key_values']
					a__        =     {
					    'num_attention_heads': 'encoder_attention_heads',
					    'hidden_size': 'd_model',
					    'attention_probs_dropout_prob': 'attention_dropout',
					}







					def __init__(      self    ,       a=9_61_03    ,       a=40_96    ,       a=16    ,       a=40_96    ,       a=16    ,       a=16    ,       a=40_96    ,       a=16    ,       a=0.0    ,       a=0.0    ,       a=True    ,       a=True    ,       a="gelu_new"    ,       a=10_24    ,       a=0.1    ,       a=0.0    ,       a=0.0    ,       a=0.02    ,       a=2    ,       a=0.0    ,       a=True    ,       a=0    ,       a=2    ,       a=1    ,       a="block_sparse"    ,       a=64    ,       a=3    ,       a=False    ,       **a    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						vocab_size
										snake_case_						:Any								=						max_position_embeddings
										snake_case_						:List[Any]								=						d_model
										snake_case_						:Any								=						encoder_ffn_dim
										snake_case_						:str								=						encoder_layers
										snake_case_						:Union[str, Any]								=						encoder_attention_heads
										snake_case_						:Any								=						decoder_ffn_dim
										snake_case_						:Dict								=						decoder_layers
										snake_case_						:List[Any]								=						decoder_attention_heads
										snake_case_						:Any								=						dropout
										snake_case_						:Union[str, Any]								=						attention_dropout
										snake_case_						:Optional[int]								=						activation_dropout
										snake_case_						:int								=						activation_function
										snake_case_						:int								=						init_std
										snake_case_						:Union[str, Any]								=						encoder_layerdrop
										snake_case_						:Dict								=						decoder_layerdrop
										snake_case_						:Any								=						classifier_dropout
										snake_case_						:Any								=						use_cache
										snake_case_						:Any								=						encoder_layers
										snake_case_						:int								=						scale_embedding  # scale factor will be sqrt(d_model) if True

										# extra config
										snake_case_						:Optional[int]								=						attention_type
										snake_case_						:int								=						block_size
										snake_case_						:Optional[int]								=						num_random_blocks
										snake_case_						:List[str]								=						use_bias

										super().__init__(
										    pad_token_id=a    ,       bos_token_id=a    ,       eos_token_id=a    ,       is_encoder_decoder=a    ,       decoder_start_token_id=a    ,       **a    ,       )



class 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					@property
					def 		_a (      self   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:int								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															    ]   )

															if self.use_past:
																				snake_case_						:int								=						{0: "batch"}
																				snake_case_						:List[str]								=						{0: "batch", 1: "past_decoder_sequence + sequence"}
															else:
																				snake_case_						:Optional[int]								=						{0: "batch", 1: "decoder_sequence"}
																				snake_case_						:Optional[Any]								=						{0: "batch", 1: "decoder_sequence"}

															if self.use_past:
																				self.fill_with_past_key_values_(a    ,       direction="inputs"   )
										elif self.task == "causal-lm":
															# TODO: figure this case out.
															snake_case_						:List[str]								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															    ]   )
															if self.use_past:
																				snake_case_	,  snake_case_						:Tuple								=						self.num_layers
																				for i in range(a   ):
																									snake_case_						:int								=						{0: "batch", 2: "past_sequence + sequence"}
																									snake_case_						:List[Any]								=						{0: "batch", 2: "past_sequence + sequence"}
										else:
															snake_case_						:List[Any]								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
															        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
															    ]   )

										return common_inputs







					@property
					def 		_a (      self   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:Optional[int]								=						super().outputs
										else:
															snake_case_						:List[str]								=						super(a    ,       self   ).outputs
															if self.use_past:
																				snake_case_	,  snake_case_						:int								=						self.num_layers
																				for i in range(a   ):
																									snake_case_						:Optional[Any]								=						{0: "batch", 2: "past_sequence + sequence"}
																									snake_case_						:Tuple								=						{0: "batch", 2: "past_sequence + sequence"}
										return common_outputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:int								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )

										# Generate decoder inputs
										snake_case_						:List[str]								=						seq_length if not self.use_past else 1
										snake_case_						:str								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )
										snake_case_						:Union[str, Any]								=						{F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
										snake_case_						:Tuple								=						dict(**a    ,       **a   )

										if self.use_past:
															if not is_torch_available():
																				raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
															else:
																				import torch
															snake_case_	,  snake_case_						:Optional[int]								=						common_inputs["input_ids"].shape
															snake_case_						:Optional[int]								=						common_inputs["decoder_input_ids"].shape[1]
															snake_case_	,  snake_case_						:Optional[int]								=						self.num_attention_heads
															snake_case_						:str								=						(
															    batch,
															    num_encoder_attention_heads,
															    encoder_seq_length,
															    self._config.hidden_size // num_encoder_attention_heads,
															)
															snake_case_						:Any								=						decoder_seq_length + 3
															snake_case_						:int								=						(
															    batch,
															    num_decoder_attention_heads,
															    decoder_past_length,
															    self._config.hidden_size // num_decoder_attention_heads,
															)

															snake_case_						:Union[str, Any]								=						torch.cat(
															    [common_inputs["decoder_attention_mask"], torch.ones(a    ,       a   )]    ,       dim=1   )

															snake_case_						:Union[str, Any]								=						[]
															# If the number of encoder and decoder layers are present in the model configuration, both are considered
															snake_case_	,  snake_case_						:Tuple								=						self.num_layers
															snake_case_						:int								=						min(a    ,       a   )
															snake_case_						:Optional[Any]								=						max(a    ,       a   ) - min_num_layers
															snake_case_						:List[Any]								=						"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

															for _ in range(a   ):
																				common_inputs["past_key_values"].append(
																				    (
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				    )   )
															# TODO: test this.
															snake_case_						:Any								=						encoder_shape if remaining_side_name == "encoder" else decoder_shape
															for _ in range(a    ,       a   ):
																				common_inputs["past_key_values"].append((torch.zeros(a   ), torch.zeros(a   ))   )
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:List[str]								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )

										if self.use_past:
															if not is_torch_available():
																				raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
															else:
																				import torch
															snake_case_	,  snake_case_						:Union[str, Any]								=						common_inputs["input_ids"].shape
															# Not using the same length for past_key_values
															snake_case_						:str								=						seqlen + 2
															snake_case_	,  snake_case_						:Union[str, Any]								=						self.num_layers
															snake_case_	,  snake_case_						:Any								=						self.num_attention_heads
															snake_case_						:List[str]								=						(
															    batch,
															    num_encoder_attention_heads,
															    past_key_values_length,
															    self._config.hidden_size // num_encoder_attention_heads,
															)

															snake_case_						:Optional[Any]								=						common_inputs["attention_mask"].dtype
															snake_case_						:Tuple								=						torch.cat(
															    [common_inputs["attention_mask"], torch.ones(a    ,       a    ,       dtype=a   )]    ,       dim=1   )
															snake_case_						:Optional[int]								=						[
															    (torch.zeros(a   ), torch.zeros(a   )) for _ in range(a   )
															]
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						compute_effective_axis_dimension(
										    a    ,       fixed_dimension=OnnxConfig.default_fixed_batch    ,       num_token_to_add=0   )

										# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
										snake_case_						:List[Any]								=						tokenizer.num_special_tokens_to_add(a   )
										snake_case_						:Tuple								=						compute_effective_axis_dimension(
										    a    ,       fixed_dimension=OnnxConfig.default_fixed_sequence    ,       num_token_to_add=a   )

										# Generate dummy inputs according to compute batch and sequence
										snake_case_						:Any								=						[" ".join([tokenizer.unk_token]   ) * seq_length] * batch_size
										snake_case_						:List[str]								=						dict(tokenizer(a    ,       return_tensors=a   )   )
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:Dict								=						self._generate_dummy_inputs_for_default_and_seqaseq_lm(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )

										elif self.task == "causal-lm":
															snake_case_						:Optional[Any]								=						self._generate_dummy_inputs_for_causal_lm(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )
										else:
															snake_case_						:List[Any]								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )

										return common_inputs







					def 		_a (      self    ,       a    ,       a    ,       a    ,       a   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:str								=						super()._flatten_past_key_values_(a    ,       a    ,       a    ,       a   )
										else:
															snake_case_						:Any								=						super(a    ,       self   )._flatten_past_key_values_(
															    a    ,       a    ,       a    ,       a   )

