


'''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


__UpperCAmelCase       =							logging.getLogger(__name__)
torch.set_grad_enabled(False)
__UpperCAmelCase       =							"cuda" if torch.cuda.is_available() else "cpu"


def  _snake_case       (					A			,   A=100			,   A=" "	)						->   Union[str, Any]:
							lowerCAmelCase__	       =     text.split(lowerCAmelCase__	)
							return [character.join(text[i : i + n]	).strip() for i in range(0			,   len(lowerCAmelCase__	)			,   lowerCAmelCase__	)]


def  _snake_case       (					A	)						->   Tuple:
							lowerCAmelCase__       , lowerCAmelCase__	       =     [], []
							for title, text in zip(documents['''title''']			,   documents['''text''']	):
														if text is not None:
																					for passage in split_text(lowerCAmelCase__	):
																												titles.append(title if title is not None else ''''''	)
																												texts.append(lowerCAmelCase__	)
							return {"title": titles, "text": texts}


def  _snake_case       (					A			,   A			,   A	)						->   Dict:
							lowerCAmelCase__	       =     ctx_tokenizer(
							    documents['''title''']			,   documents['''text''']			,   truncation=lowerCAmelCase__			,   padding='''longest'''			,   return_tensors='''pt'''	)['''input_ids''']
							lowerCAmelCase__	       =     ctx_encoder(input_ids.to(device=lowerCAmelCase__	)			,   return_dict=lowerCAmelCase__	).pooler_output
							return {"embeddings": embeddings.detach().cpu().numpy()}


def  _snake_case       (					A			,   A			,   A			,   )						->   List[str]:
							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
							lowerCAmelCase__	       =     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
							lowerCAmelCase__	       =     dataset.map(lowerCAmelCase__			,   batched=lowerCAmelCase__			,   num_proc=processing_args.num_proc	)

							# And compute the embeddings
							lowerCAmelCase__	       =     DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	).to(device=lowerCAmelCase__	)
							lowerCAmelCase__	       =     DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	)
							lowerCAmelCase__	       =     Features(
							    {'''text''': Value('''string'''	), '''title''': Value('''string'''	), '''embeddings''': Sequence(Value('''float32'''	)	)}	)  # optional, save as float32 instead of float64 to save space
							lowerCAmelCase__	       =     dataset.map(
							    partial(lowerCAmelCase__			,   ctx_encoder=lowerCAmelCase__			,   ctx_tokenizer=lowerCAmelCase__	)			,   batched=lowerCAmelCase__			,   batch_size=processing_args.batch_size			,   features=lowerCAmelCase__			,   )

							# And finally save your dataset
							lowerCAmelCase__	       =     os.path.join(rag_example_args.output_dir			,   '''my_knowledge_dataset'''	)
							dataset.save_to_disk(lowerCAmelCase__	)
							# 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
							lowerCAmelCase__	       =     faiss.IndexHNSWFlat(index_hnsw_args.d			,   index_hnsw_args.m			,   faiss.METRIC_INNER_PRODUCT	)
							dataset.add_faiss_index('''embeddings'''			,   custom_index=lowerCAmelCase__	)

							# And save the index
							lowerCAmelCase__	       =     os.path.join(rag_example_args.output_dir			,   '''my_knowledge_dataset_hnsw_index.faiss'''	)
							dataset.get_index('''embeddings'''	).save(lowerCAmelCase__	)
							# dataset.load_faiss_index("embeddings", index_path)  # to reload the index




@dataclass
class 		a__    :
							lowercase__      :       str     						=				field(
							    default=str(Path(SCREAMING_SNAKE_CASE__  ).parent / "test_run" / "dummy-kb" / "my_knowledge_dataset.csv"  )    ,    metadata={"help": "Path to a tab-separated csv file with columns \'title\' and \'text\'"}    ,    )
							lowercase__      :       Optional[str]     						=				field(
							    default=SCREAMING_SNAKE_CASE__    ,    metadata={"help": "Question that is passed as input to RAG. Default is \'What does Moses\' rod turn into ?\'."}    ,    )
							lowercase__      :       str     						=				field(
							    default="facebook/rag-sequence-nq"    ,    metadata={"help": "The RAG model to use. Either \'facebook/rag-sequence-nq\' or \'facebook/rag-token-nq\'"}    ,    )
							lowercase__      :       str     						=				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\'"
							        )
							    }    ,    )
							lowercase__      :       Optional[str]     						=				field(
							    default=str(Path(SCREAMING_SNAKE_CASE__  ).parent / "test_run" / "dummy-kb"  )    ,    metadata={"help": "Path to a directory where the dataset passages and the index will be saved"}    ,    )




@dataclass
class 		a__    :
							lowercase__      :       Optional[int]     						=				field(
							    default=SCREAMING_SNAKE_CASE__    ,    metadata={
							        "help": "The number of processes to use to split the documents into passages. Default is single process."
							    }    ,    )
							lowercase__      :       int     						=				field(
							    default=1_6    ,    metadata={
							        "help": "The batch size to use when computing the passages embeddings using the DPR context encoder."
							    }    ,    )




@dataclass
class 		a__    :
							lowercase__      :       int     						=				field(
							    default=7_6_8    ,    metadata={"help": "The dimension of the embeddings to pass to the HNSW Faiss index."}    ,    )
							lowercase__      :       int     						=				field(
							    default=1_2_8    ,    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       =							HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
		__UpperCAmelCase       =							parser.parse_args_into_dataclasses()
		with TemporaryDirectory() as tmp_dir:
				__UpperCAmelCase       =							rag_example_args.output_dir or tmp_dir
				main(rag_example_args, processing_args, index_hnsw_args)



'''simple docstring'''






import json
from typing import List, Optional, Tuple

from tokenizers import normalizers
from tokenizers.pre_tokenizers import BertPreTokenizer, PreTokenizer

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_roformer import RoFormerTokenizer
from .tokenization_utils import JiebaPreTokenizer


__UpperCAmelCase       =							logging.get_logger(__name__)

__UpperCAmelCase       =							{'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

__UpperCAmelCase       =							{
    '''vocab_file''': {
        '''junnyu/roformer_chinese_small''': '''https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/vocab.txt''',
        '''junnyu/roformer_chinese_base''': '''https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/vocab.txt''',
        '''junnyu/roformer_chinese_char_small''': (
            '''https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/vocab.txt'''
        ),
        '''junnyu/roformer_chinese_char_base''': (
            '''https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/vocab.txt'''
        ),
        '''junnyu/roformer_small_discriminator''': (
            '''https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/vocab.txt'''
        ),
        '''junnyu/roformer_small_generator''': (
            '''https://huggingface.co/junnyu/roformer_small_generator/resolve/main/vocab.txt'''
        ),
    }
}

__UpperCAmelCase       =							{
    '''junnyu/roformer_chinese_small''': 1_536,
    '''junnyu/roformer_chinese_base''': 1_536,
    '''junnyu/roformer_chinese_char_small''': 512,
    '''junnyu/roformer_chinese_char_base''': 512,
    '''junnyu/roformer_small_discriminator''': 128,
    '''junnyu/roformer_small_generator''': 128,
}


__UpperCAmelCase       =							{
    '''junnyu/roformer_chinese_small''': {'''do_lower_case''': True},
    '''junnyu/roformer_chinese_base''': {'''do_lower_case''': True},
    '''junnyu/roformer_chinese_char_small''': {'''do_lower_case''': True},
    '''junnyu/roformer_chinese_char_base''': {'''do_lower_case''': True},
    '''junnyu/roformer_small_discriminator''': {'''do_lower_case''': True},
    '''junnyu/roformer_small_generator''': {'''do_lower_case''': True},
}




class 		a__    (							a__  ):





							'''simple docstring'''






							lowercase__      :       int     						=				VOCAB_FILES_NAMES
							lowercase__      :       Optional[int]     						=				PRETRAINED_VOCAB_FILES_MAP
							lowercase__      :       Any     						=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowercase__      :       Union[str, Any]     						=				PRETRAINED_INIT_CONFIGURATION
							lowercase__      :       Tuple     						=				RoFormerTokenizer







							def __init__(		self ,   lowerCamelCase_=None ,   lowerCamelCase_=None ,   lowerCamelCase_=True ,   lowerCamelCase_="[UNK]" ,   lowerCamelCase_="[SEP]" ,   lowerCamelCase_="[PAD]" ,   lowerCamelCase_="[CLS]" ,   lowerCamelCase_="[MASK]" ,   lowerCamelCase_=True ,   lowerCamelCase_=None ,   **lowerCamelCase_ ,   )	->							Tuple:
														super().__init__(
														    lowerCamelCase_ ,   tokenizer_file=lowerCamelCase_ ,   do_lower_case=lowerCamelCase_ ,   unk_token=lowerCamelCase_ ,   sep_token=lowerCamelCase_ ,   pad_token=lowerCamelCase_ ,   cls_token=lowerCamelCase_ ,   mask_token=lowerCamelCase_ ,   tokenize_chinese_chars=lowerCamelCase_ ,   strip_accents=lowerCamelCase_ ,   **lowerCamelCase_ ,   )

														lowerCAmelCase__	       =     json.loads(self.backend_tokenizer.normalizer.__getstate__()       )
														if (
														    pre_tok_state.get('''lowercase''' ,   lowerCamelCase_       ) != do_lower_case
														    or pre_tok_state.get('''strip_accents''' ,   lowerCamelCase_       ) != strip_accents
														):
																					lowerCAmelCase__	       =     getattr(lowerCamelCase_ ,   pre_tok_state.pop('''type'''       )       )
																					lowerCAmelCase__	       =     do_lower_case
																					lowerCAmelCase__	       =     strip_accents
																					lowerCAmelCase__	       =     pre_tok_class(**lowerCamelCase_       )

														lowerCAmelCase__	       =     do_lower_case







							def __getstate__(		self       )	->							Any:
														lowerCAmelCase__	       =     self.__dict__.copy()
														lowerCAmelCase__	       =     BertPreTokenizer()
														return state







							def __setstate__(		self ,   lowerCamelCase_       )	->							List[Any]:
														lowerCAmelCase__	       =     d
														lowerCAmelCase__	       =     self.__dict__['''_tokenizer'''].get_vocab()
														lowerCAmelCase__	       =     PreTokenizer.custom(JiebaPreTokenizer(lowerCamelCase_       )       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_=None       )	->							Union[str, Any]:
														lowerCAmelCase__	       =     [self.cls_token_id] + token_ids_a + [self.sep_token_id]

														if token_ids_a:
																					output += token_ids_a + [self.sep_token_id]

														return output







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ = None       )	->							List[int]:
														lowerCAmelCase__	       =     [self.sep_token_id]
														lowerCAmelCase__	       =     [self.cls_token_id]
														if token_ids_a is None:
																					return len(cls + token_ids_a + sep       ) * [0]
														return len(cls + token_ids_a + sep       ) * [0] + len(token_ids_a + sep       ) * [1]







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_ = None       )	->							Tuple[str]:
														lowerCAmelCase__	       =     self._tokenizer.model.save(lowerCamelCase_ ,   name=lowerCamelCase_       )
														return tuple(lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_ ,   lowerCamelCase_=None ,   lowerCamelCase_=None ,   lowerCamelCase_=False ,   **lowerCamelCase_ ,   )	->							Union[str, Any]:
														lowerCAmelCase__	       =     BertPreTokenizer()
														return super().save_pretrained(lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   lowerCamelCase_ ,   **lowerCamelCase_       )