



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


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





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase=100     ,   _lowerCAmelCase=" "		)		-> List[str]:
      snake_case__  :    Dict    						=      text.split(snake_case_		)
      return [character.join(text[i : i + n]		).strip() for i in range(0     ,   len(snake_case_		)     ,   snake_case_		)]





def 						__snake_case(				_lowerCAmelCase		)		-> dict:
      snake_case__  :    Optional[int]    						=      [], []
      for title, text in zip(documents["""title"""]     ,   documents["""text"""]		):
            if text is not None:
                  for passage in split_text(snake_case_		):
                        titles.append(title if title is not None else """"""		)
                        texts.append(snake_case_		)
      return {"title": titles, "text": texts}





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase		)		-> dict:
      snake_case__  :    Tuple    						=      ctx_tokenizer(
          documents["""title"""]     ,   documents["""text"""]     ,   truncation=snake_case_     ,   padding="""longest"""     ,   return_tensors="""pt"""		)['''input_ids''']
      snake_case__  :    int    						=      ctx_encoder(input_ids.to(device=snake_case_		)     ,   return_dict=snake_case_		).pooler_output
      return {"embeddings": embeddings.detach().cpu().numpy()}





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase     ,   )		-> Tuple:
      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__  :    Optional[Any]    						=      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__  :    str    						=      dataset.map(snake_case_     ,   batched=snake_case_     ,   num_proc=processing_args.num_proc		)

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

      # And finally save your dataset
      snake_case__  :    int    						=      os.path.join(rag_example_args.output_dir     ,   """my_knowledge_dataset"""		)
      dataset.save_to_disk(snake_case_		)
      # 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__  :    Tuple    						=      faiss.IndexHNSWFlat(index_hnsw_args.d     ,   index_hnsw_args.m     ,   faiss.METRIC_INNER_PRODUCT		)
      dataset.add_faiss_index("""embeddings"""     ,   custom_index=snake_case_		)

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







@dataclass
class     UpperCAmelCase_       :
    """simple docstring"""







    lowercase					=		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					=		field(
        default=SCREAMING_SNAKE_CASE__		,		metadata={"help": "Question that is passed as input to RAG. Default is 'What does Moses' rod turn into ?'."}		,		)
    lowercase					=		field(
        default="facebook/rag-sequence-nq"		,		metadata={"help": "The RAG model to use. Either 'facebook/rag-sequence-nq' or 'facebook/rag-token-nq'"}		,		)
    lowercase					=		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					=		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     UpperCAmelCase_       :
    """simple docstring"""







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







@dataclass
class     UpperCAmelCase_       :
    """simple docstring"""







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

       __a          =		HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
       __a       ,						__a       ,						__a          =		parser.parse_args_into_dataclasses()
       with TemporaryDirectory() as tmp_dir:
              __a          =		rag_example_args.output_dir or tmp_dir
              main(rag_example_args, processing_args, index_hnsw_args)







"""simple docstring"""



import uuid
from typing import Any, Dict, List, Optional, Union

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
       import tensorflow as tf

if is_torch_available():
       import torch


a_ =			logging.get_logger(__name__)






class 		__snake_case					:





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase = None					,						__lowerCamelCase = None					,						__lowerCamelCase=None					,						__lowerCamelCase=None  ):



     '''simple docstring'''




     if not conversation_id:
         __A    :						List[Any]						 =     uuid.uuida()
     if past_user_inputs is None:
         __A    :						List[str]						 =     []
     if generated_responses is None:
         __A    :						Tuple						 =     []

     __A    :						uuid.UUID						 =     conversation_id
     __A    :						List[str]						 =     past_user_inputs
     __A    :						List[str]						 =     generated_responses
     __A    :						Optional[str]						 =     text

 def __eq__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     if not isinstance(__lowerCamelCase					,						__lowerCamelCase  ):
         return False
     if self.uuid == other.uuid:
         return True
     return (
         self.new_user_input == other.new_user_input
         and self.past_user_inputs == other.past_user_inputs
         and self.generated_responses == other.generated_responses
     )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase = False  ):



     '''simple docstring'''




     if self.new_user_input:
         if overwrite:
             logger.warning(
                 F"""User input added while unprocessed input was existing: \"{self.new_user_input}\" was overwritten """
                 F"""with: \"{text}\"."""  )
             __A    :						str						 =     text
         else:
             logger.warning(
                 F"""User input added while unprocessed input was existing: \"{self.new_user_input}\" new input """
                 F"""ignored: \"{text}\". Set `overwrite` to True to overwrite unprocessed user input"""  )
     else:
         __A    :						Union[str, Any]						 =     text

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     if self.new_user_input:
         self.past_user_inputs.append(self.new_user_input  )
     __A    :						List[Any]						 =     None

 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     self.generated_responses.append(__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     for user_input, generated_response in zip(self.past_user_inputs					,						self.generated_responses  ):
         yield True, user_input
         yield False, generated_response
     if self.new_user_input:
         yield True, self.new_user_input







 def __repr__(    self  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     F"""Conversation id: {self.uuid} \n"""
     for is_user, text in self.iter_texts():
         __A    :						Tuple						 =     '''user''' if is_user else '''bot'''
         output += F"""{name} >> {text} \n"""
     return output






@add_end_docstrings(
    SCREAMING_SNAKE_CASE__							,							R"""
        min_length_for_response (`int`, *optional*, defaults to 32):
            The minimum length (in number of tokens) for a response.
        minimum_tokens (`int`, *optional*, defaults to 10):
            The minimum length of tokens to leave for a response.
    """							,							)
class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 def __init__(    self					,						*__lowerCamelCase					,						**__lowerCamelCase  ):



     '''simple docstring'''




     super().__init__(*__lowerCamelCase					,						**__lowerCamelCase  )
     if self.tokenizer.pad_token_id is None:
         __A    :						Union[str, Any]						 =     self.tokenizer.eos_token

 def 					UpperCamelCase__(    self					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						**__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						str						 =     {}
     __A    :						List[str]						 =     {}
     __A    :						Any						 =     {}

     if min_length_for_response is not None:
         __A    :						int						 =     min_length_for_response
     if minimum_tokens is not None:
         __A    :						Any						 =     minimum_tokens

     if "max_length" in generate_kwargs:
         __A    :						List[Any]						 =     generate_kwargs['''max_length''']
         # self.max_length = generate_kwargs.get("max_length", self.model.config.max_length)
     if clean_up_tokenization_spaces is not None:
         __A    :						str						 =     clean_up_tokenization_spaces

     if generate_kwargs:
         forward_params.update(__lowerCamelCase  )
     return preprocess_params, forward_params, postprocess_params

 def __call__(    self					,						__lowerCamelCase					,						__lowerCamelCase=0					,						**__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Any						 =     super().__call__(__lowerCamelCase					,						num_workers=__lowerCamelCase					,						**__lowerCamelCase  )
     if isinstance(__lowerCamelCase					,						__lowerCamelCase  ) and len(__lowerCamelCase  ) == 1:
         return outputs[0]
     return outputs

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase=32  ):



     '''simple docstring'''




     if not isinstance(__lowerCamelCase					,						__lowerCamelCase  ):
         raise ValueError('''ConversationalPipeline, expects Conversation as inputs'''  )
     if conversation.new_user_input is None:
         raise ValueError(
             F"""Conversation with UUID {type(conversation.uuid  )} does not contain new user input to process. """
             '''Add user inputs with the conversation\'s `add_user_input` method'''  )
     if hasattr(self.tokenizer					,						'''_build_conversation_input_ids'''  ):
         __A    :						List[Any]						 =     self.tokenizer._build_conversation_input_ids(__lowerCamelCase  )
     else:
         # If the tokenizer cannot handle conversations, we default to only the old version
         __A    :						int						 =     self._legacy_parse_and_tokenize(__lowerCamelCase  )

     if self.framework == "pt":
         __A    :						Union[str, Any]						 =     torch.LongTensor([input_ids]  )
     elif self.framework == "tf":
         __A    :						int						 =     tf.constant([input_ids]  )
     return {"input_ids": input_ids, "conversation": conversation}

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase=10					,						**__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Tuple						 =     generate_kwargs.get('''max_length'''					,						self.model.config.max_length  )

     __A    :						str						 =     model_inputs['''input_ids'''].shape[1]
     if max_length - minimum_tokens < n:
         logger.warning(F"""Conversation input is to long ({n}), trimming it to ({max_length} - {minimum_tokens})"""  )
         __A    :						str						 =     max_length - minimum_tokens
         __A    :						Any						 =     model_inputs['''input_ids'''][:, -trim:]
         if "attention_mask" in model_inputs:
             __A    :						Union[str, Any]						 =     model_inputs['''attention_mask'''][:, -trim:]
     __A    :						Dict						 =     model_inputs.pop('''conversation'''  )
     __A    :						List[str]						 =     max_length
     __A    :						Dict						 =     self.model.generate(**__lowerCamelCase					,						**__lowerCamelCase  )
     if self.model.config.is_encoder_decoder:
         __A    :						Any						 =     1
     else:
         __A    :						List[Any]						 =     n
     return {"output_ids": output_ids[:, start_position:], "conversation": conversation}

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase=True  ):



     '''simple docstring'''




     __A    :						int						 =     model_outputs['''output_ids''']
     __A    :						Optional[int]						 =     self.tokenizer.decode(
         output_ids[0]					,						skip_special_tokens=__lowerCamelCase					,						clean_up_tokenization_spaces=__lowerCamelCase					,						)
     __A    :						Dict						 =     model_outputs['''conversation''']
     conversation.mark_processed()
     conversation.append_response(__lowerCamelCase  )
     return conversation







 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     __A    :						Tuple						 =     self.tokenizer.eos_token_id
     __A    :						List[str]						 =     []
     for is_user, text in conversation.iter_texts():
         if eos_token_id is not None:
             input_ids.extend(self.tokenizer.encode(__lowerCamelCase					,						add_special_tokens=__lowerCamelCase  ) + [eos_token_id]  )
         else:
             input_ids.extend(self.tokenizer.encode(__lowerCamelCase					,						add_special_tokens=__lowerCamelCase  )  )

     if len(__lowerCamelCase  ) > self.tokenizer.model_max_length:
         __A    :						List[str]						 =     input_ids[-self.tokenizer.model_max_length :]
     return input_ids




