


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


A__		: List[str]												=  logging.getLogger(__name__)
torch.set_grad_enabled(False)
A__		: List[Any]												=  """cuda""" if torch.cuda.is_available() else """cpu"""



def       UpperCAmelCase__ (     UpperCAmelCase_   :    str				,   UpperCAmelCase_   :    Optional[Any]=1_00				,   UpperCAmelCase_   :    Union[str, Any]=" "							)					->		List[str]:
     __lowerCamelCase       :     Optional[Any]        =		text.split(UpperCAmelCase_							)
     return [character.join(text[i : i + n]							).strip() for i in range(0				,   len(UpperCAmelCase_							)				,   UpperCAmelCase_							)]



def       UpperCAmelCase__ (     UpperCAmelCase_   :    dict							)					->		dict:
     __lowerCamelCase    ,	__lowerCamelCase       :     Any        =		[], []
     for title, text in zip(documents['title']				,   documents['text']							):
          if text is not None:
               for passage in split_text(UpperCAmelCase_							):
                    titles.append(title if title is not None else ''							)
                    texts.append(UpperCAmelCase_							)
     return {"title": titles, "text": texts}



def       UpperCAmelCase__ (     UpperCAmelCase_   :    dict				,   UpperCAmelCase_   :    DPRContextEncoder				,   UpperCAmelCase_   :    DPRContextEncoderTokenizerFast							)					->		dict:
     __lowerCamelCase       :     Tuple        =		ctx_tokenizer(
         documents['title']				,   documents['text']				,   truncation=UpperCAmelCase_				,   padding='longest'				,   return_tensors='pt'							)['input_ids']
     __lowerCamelCase       :     List[Any]        =		ctx_encoder(input_ids.to(device=UpperCAmelCase_							)				,   return_dict=UpperCAmelCase_							).pooler_output
     return {"embeddings": embeddings.detach().cpu().numpy()}



def       UpperCAmelCase__ (     UpperCAmelCase_   :    "RagExampleArguments"				,   UpperCAmelCase_   :    "ProcessingArguments"				,   UpperCAmelCase_   :    "IndexHnswArguments"				,   )					->		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       :     str        =		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       :     Dict        =		dataset.map(UpperCAmelCase_				,   batched=UpperCAmelCase_				,   num_proc=processing_args.num_proc							)

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

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

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

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

     # Easy way to load the model
     __lowerCamelCase       :     Union[str, Any]        =		RagRetriever.from_pretrained(
         rag_example_args.rag_model_name				,   index_name='custom'				,   indexed_dataset=UpperCAmelCase_							)
     __lowerCamelCase       :     Tuple        =		RagSequenceForGeneration.from_pretrained(rag_example_args.rag_model_name				,   retriever=UpperCAmelCase_							)
     __lowerCamelCase       :     Union[str, 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'							)
     ######################################

     __lowerCamelCase       :     int        =		rag_example_args.question or 'What does Moses\' rod turn into ?'
     __lowerCamelCase       :     Optional[Any]        =		tokenizer.question_encoder(UpperCAmelCase_				,   return_tensors='pt'							)['input_ids']
     __lowerCamelCase       :     Union[str, Any]        =		model.generate(UpperCAmelCase_							)
     __lowerCamelCase       :     str        =		tokenizer.batch_decode(UpperCAmelCase_				,   skip_special_tokens=UpperCAmelCase_							)[0]
     logger.info('Q: ' + question							)
     logger.info('A: ' + generated_string							)





@dataclass
class        UpperCAmelCase_  :





    """simple docstring"""






    lowerCamelCase  : str												=      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\''} ,		)
    lowerCamelCase  : Optional[str]												=      field(
        default=_UpperCAmelCase ,		metadata={'help': 'Question that is passed as input to RAG. Default is \'What does Moses\' rod turn into ?\'.'} ,		)
    lowerCamelCase  : str												=      field(
        default='facebook/rag-sequence-nq' ,		metadata={'help': 'The RAG model to use. Either \'facebook/rag-sequence-nq\' or \'facebook/rag-token-nq\''} ,		)
    lowerCamelCase  : 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\''
            )
        } ,		)
    lowerCamelCase  : Optional[str]												=      field(
        default=_UpperCAmelCase ,		metadata={'help': 'Path to a directory where the dataset passages and the index will be saved'} ,		)







@dataclass
class        UpperCAmelCase_  :





    """simple docstring"""






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







@dataclass
class        UpperCAmelCase_  :





    """simple docstring"""






    lowerCamelCase  : int												=      field(
        default=7_6_8 ,		metadata={'help': 'The dimension of the embeddings to pass to the HNSW Faiss index.'} ,		)
    lowerCamelCase  : 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)

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










'''simple docstring'''



from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_clip import CLIPTokenizer


A__		: Optional[int]												=  logging.get_logger(__name__)

A__		: Optional[Any]												=  {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt""", """tokenizer_file""": """tokenizer.json"""}

A__		: Optional[int]												=  {
    """vocab_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/vocab.json""",
    },
    """merges_file""": {
        """openai/clip-vit-base-patch32""": """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/merges.txt""",
    },
    """tokenizer_file""": {
        """openai/clip-vit-base-patch32""": (
            """https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/tokenizer.json"""
        ),
    },
}

A__		: Any												=  {
    """openai/clip-vit-base-patch32""": 77,
}





class        UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    lowerCamelCase  : Union[str, Any]												=      VOCAB_FILES_NAMES
    lowerCamelCase  : Optional[int]												=      PRETRAINED_VOCAB_FILES_MAP
    lowerCamelCase  : str												=      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    lowerCamelCase  : Union[str, Any]												=      ['input_ids', 'attention_mask']
    lowerCamelCase  : Optional[Any]												=      CLIPTokenizer







    def __init__(							self      ,					SCREAMING_SNAKE_CASE_=None      ,					SCREAMING_SNAKE_CASE_=None      ,					SCREAMING_SNAKE_CASE_=None      ,					SCREAMING_SNAKE_CASE_="<|endoftext|>"      ,					SCREAMING_SNAKE_CASE_="<|startoftext|>"      ,					SCREAMING_SNAKE_CASE_="<|endoftext|>"      ,					SCREAMING_SNAKE_CASE_="<|endoftext|>"      ,					**SCREAMING_SNAKE_CASE_      ,					)    ->					Union[str, Any]:
         super().__init__(
             SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_      ,					tokenizer_file=SCREAMING_SNAKE_CASE_      ,					unk_token=SCREAMING_SNAKE_CASE_      ,					bos_token=SCREAMING_SNAKE_CASE_      ,					eos_token=SCREAMING_SNAKE_CASE_      ,					pad_token=SCREAMING_SNAKE_CASE_      ,					**SCREAMING_SNAKE_CASE_      ,					)

         if not isinstance(self.backend_tokenizer.pre_tokenizer      ,					pre_tokenizers.Sequence     ):
              raise ValueError(
                  'The `backend_tokenizer` provided does not match the expected format. The CLIP tokenizer has been'
                  ' heavily modified from transformers version 4.17.0. You need to convert the tokenizer you are using'
                  ' to be compatible with this version.The easiest way to do so is'
                  ' `CLIPTokenizerFast.from_pretrained("path_to_local_folder_or_hub_repo, from_slow=True)`. If you want'
                  ' to use your existing tokenizer, you will have to revert to a version prior to 4.17.0 of'
                  ' transformers.'     )

         self._wrap_decode_method_backend_tokenizer()







    def 				lowercase_    (							self     )    ->					Optional[Any]:
         __lowerCamelCase       :     Union[str, Any]        =		self.backend_tokenizer.decode

         def new_decode_method(*SCREAMING_SNAKE_CASE_      ,					**SCREAMING_SNAKE_CASE_     ):
              __lowerCamelCase       :     int        =		orig_decode_method(*SCREAMING_SNAKE_CASE_      ,					**SCREAMING_SNAKE_CASE_     )
              __lowerCamelCase       :     List[str]        =		text.replace(self.backend_tokenizer.model.end_of_word_suffix      ,					' '     ).strip()
              return text

         __lowerCamelCase       :     Union[str, Any]        =		new_decode_method







    def 				lowercase_    (							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_ = None     )    ->					List[int]:
         __lowerCamelCase       :     Optional[int]        =		[self.bos_token_id]
         __lowerCamelCase       :     Tuple        =		[self.eos_token_id]

         if token_ids_a is None:
              return bos_token + token_ids_a + eos_token
         return bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token







    def 				lowercase_    (							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_ = None     )    ->					List[int]:
         __lowerCamelCase       :     Dict        =		[self.bos_token_id]
         __lowerCamelCase       :     Optional[Any]        =		[self.eos_token_id]

         if token_ids_a is None:
              return len(bos_token + token_ids_a + eos_token     ) * [0]
         return len(bos_token + token_ids_a + eos_token + eos_token + token_ids_a + eos_token     ) * [0]







    def 				lowercase_    (							self      ,					SCREAMING_SNAKE_CASE_      ,					SCREAMING_SNAKE_CASE_ = None     )    ->					Tuple[str]:
         __lowerCamelCase       :     int        =		self._tokenizer.model.save(SCREAMING_SNAKE_CASE_      ,					name=SCREAMING_SNAKE_CASE_     )
         return tuple(SCREAMING_SNAKE_CASE_     )






