



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 		UpperCamelCase	(   snake_case__       :    str       ,     snake_case__       :    Optional[int]=100       ,     snake_case__       :    Any=" "	)	->     List[str]:
   UpperCamelCase    :  int							=	text.split(snake_case__	)
   return [character.join(text[i : i + n]	).strip() for i in range(0       ,     len(snake_case__	)       ,     snake_case__	)]
def 		UpperCamelCase	(   snake_case__       :    dict	)	->     dict:
   UpperCamelCase	,     UpperCamelCase    :  Union[str, Any]							=	[], []
   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 		UpperCamelCase	(   snake_case__       :    dict       ,     snake_case__       :    DPRContextEncoder       ,     snake_case__       :    DPRContextEncoderTokenizerFast	)	->     dict:
   UpperCamelCase    :  Optional[Any]							=	ctx_tokenizer(
       documents['title']       ,     documents['text']       ,     truncation=snake_case__       ,     padding='longest'       ,     return_tensors='pt'	)['input_ids']
   UpperCamelCase    :  List[str]							=	ctx_encoder(input_ids.to(device=snake_case__	)       ,     return_dict=snake_case__	).pooler_output
   return {"embeddings": embeddings.detach().cpu().numpy()}
def 		UpperCamelCase	(   snake_case__       :    "RagExampleArguments"       ,     snake_case__       :    "ProcessingArguments"       ,     snake_case__       :    "IndexHnswArguments"       ,     )	->     Any:
   ######################################
   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
   UpperCamelCase    :  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
   UpperCamelCase    :  int							=	dataset.map(snake_case__       ,     batched=snake_case__       ,     num_proc=processing_args.num_proc	)

   # And compute the embeddings
   UpperCamelCase    :  Dict							=	DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	).to(device=snake_case__	)
   UpperCamelCase    :  int							=	DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name	)
   UpperCamelCase    :  Optional[int]							=	Features(
       {'text': Value('string'	), 'title': Value('string'	), 'embeddings': Sequence(Value('float32'	)	)}	)  # optional, save as float32 instead of float64 to save space
   UpperCamelCase    :  List[str]							=	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
   UpperCamelCase    :  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
   UpperCamelCase    :  str							=	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
   UpperCamelCase    :  Optional[int]							=	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 				lowerCAmelCase_		:
       UpperCAmelCase__				:					str        =  field(
           default=str(Path(a__							).parent / "test_run" / "dummy-kb" / "my_knowledge_dataset.csv"							)       ,    metadata={"help": "Path to a tab-separated csv file with columns 'title' and 'text'"}       ,    )
       UpperCAmelCase__				:					Optional[str]        =  field(
           default=a__       ,    metadata={"help": "Question that is passed as input to RAG. Default is 'What does Moses' rod turn into ?'."}       ,    )
       UpperCAmelCase__				:					str        =  field(
           default="facebook/rag-sequence-nq"       ,    metadata={"help": "The RAG model to use. Either 'facebook/rag-sequence-nq' or 'facebook/rag-token-nq'"}       ,    )
       UpperCAmelCase__				:					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'"
               )
           }       ,    )
       UpperCAmelCase__				:					Optional[str]        =  field(
           default=str(Path(a__							).parent / "test_run" / "dummy-kb"							)       ,    metadata={"help": "Path to a directory where the dataset passages and the index will be saved"}       ,    )






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






@dataclass
class 				lowerCAmelCase_		:
       UpperCAmelCase__				:					int        =  field(
           default=768       ,    metadata={"help": "The dimension of the embeddings to pass to the HNSW Faiss index."}       ,    )
       UpperCAmelCase__				:					int        =  field(
           default=128       ,    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   ,  __UpperCAmelCase   ,  __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)








import os
import sys
import warnings
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen
from ..table import array_cast
from ..utils.file_utils import is_local_path
from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
    import PIL.Image

    from .features import FeatureType


__UpperCAmelCase          =       None
__UpperCAmelCase          =       '''<''' if sys.byteorder == '''little''' else '''>'''
# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
__UpperCAmelCase          =       [
    np.dtype('''|b1'''),
    np.dtype('''|u1'''),
    np.dtype('''<u2'''),
    np.dtype('''>u2'''),
    np.dtype('''<i2'''),
    np.dtype('''>i2'''),
    np.dtype('''<u4'''),
    np.dtype('''>u4'''),
    np.dtype('''<i4'''),
    np.dtype('''>i4'''),
    np.dtype('''<f4'''),
    np.dtype('''>f4'''),
    np.dtype('''<f8'''),
    np.dtype('''>f8'''),
]



@dataclass
class 				lowerCAmelCase_		:
       UpperCAmelCase__				:					bool        =  True
       UpperCAmelCase__				:					Optional[str]        =  None
       # Automatically constructed
       UpperCAmelCase__				:					ClassVar[str]        =  "PIL.Image.Image"
       UpperCAmelCase__				:					ClassVar[Any]        =  pa.struct({"bytes": pa.binary(), "path": pa.string()}							)
       UpperCAmelCase__				:					str        =  field(default="Image"       ,    init=a__       ,    repr=a__							)
       def __call__(			self       )     ->		Any:
          return self.pa_type
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_       )     ->		dict:
          if config.PIL_AVAILABLE:
             import PIL.Image
          else:
             raise ImportError('To support encoding images, please install \'Pillow\'.'       )

          if isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             UpperCamelCase    :  Optional[int]							=	np.array(SCREAMING_SNAKE_CASE_       )

          if isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             return {"path": value, "bytes": None}
          elif isinstance(SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_       ):
             return {"path": None, "bytes": value}
          elif isinstance(SCREAMING_SNAKE_CASE_,			np.ndarray       ):
             # convert the image array to PNG/TIFF bytes
             return encode_np_array(SCREAMING_SNAKE_CASE_       )
          elif isinstance(SCREAMING_SNAKE_CASE_,			PIL.Image.Image       ):
             # convert the PIL image to bytes (default format is PNG/TIFF)
             return encode_pil_image(SCREAMING_SNAKE_CASE_       )
          elif value.get('path'       ) is not None and os.path.isfile(value['path']       ):
             # we set "bytes": None to not duplicate the data if they're already available locally
             return {"bytes": None, "path": value.get('path'       )}
          elif value.get('bytes'       ) is not None or value.get('path'       ) is not None:
             # store the image bytes, and path is used to infer the image format using the file extension
             return {"bytes": value.get('bytes'       ), "path": value.get('path'       )}
          else:
             raise ValueError(
                 F"""An image sample should have one of 'path' or 'bytes' but they are missing or None in {value}."""       )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_,			SCREAMING_SNAKE_CASE_=None       )     ->		"PIL.Image.Image":
          if not self.decode:
             raise RuntimeError('Decoding is disabled for this feature. Please use Image(decode=True) instead.'       )

          if config.PIL_AVAILABLE:
             import PIL.Image
          else:
             raise ImportError('To support decoding images, please install \'Pillow\'.'       )

          if token_per_repo_id is None:
             UpperCamelCase    :  Any							=	{}

          UpperCamelCase	,     UpperCamelCase    :  Union[str, Any]							=	value['path'], value['bytes']
          if bytes_ is None:
             if path is None:
                raise ValueError(F"""An image should have one of 'path' or 'bytes' but both are None in {value}."""       )
             else:
                if is_local_path(SCREAMING_SNAKE_CASE_       ):
                   UpperCamelCase    :  int							=	PIL.Image.open(SCREAMING_SNAKE_CASE_       )
                else:
                   UpperCamelCase    :  int							=	path.split('::'       )[-1]
                   try:
                      UpperCamelCase    :  Optional[Any]							=	string_to_dict(SCREAMING_SNAKE_CASE_,			config.HUB_DATASETS_URL       )['repo_id']
                      UpperCamelCase    :  str							=	token_per_repo_id.get(SCREAMING_SNAKE_CASE_       )
                   except ValueError:
                      UpperCamelCase    :  Tuple							=	None
                   with xopen(SCREAMING_SNAKE_CASE_,			'rb',			use_auth_token=SCREAMING_SNAKE_CASE_       ) as f:
                      UpperCamelCase    :  Optional[int]							=	BytesIO(f.read()       )
                   UpperCamelCase    :  int							=	PIL.Image.open(bytes_       )
          else:
             UpperCamelCase    :  Optional[int]							=	PIL.Image.open(BytesIO(bytes_       )       )
          image.load()  # to avoid "Too many open files" errors
          return image
       def 					snake_case_       (			self       )     ->		Union["FeatureType", Dict[str, "FeatureType"]]:
          from .features import Value

          return (
              self
              if self.decode
              else {
                  "bytes": Value('binary'       ),
                  "path": Value('string'       ),
              }
          )
       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_       )     ->		pa.StructArray:
          if pa.types.is_string(storage.type       ):
             UpperCamelCase    :  List[str]							=	pa.array([None] * len(SCREAMING_SNAKE_CASE_       ),			type=pa.binary()       )
             UpperCamelCase    :  Dict							=	pa.StructArray.from_arrays([bytes_array, storage],			['bytes', 'path'],			mask=storage.is_null()       )
          elif pa.types.is_binary(storage.type       ):
             UpperCamelCase    :  Optional[int]							=	pa.array([None] * len(SCREAMING_SNAKE_CASE_       ),			type=pa.string()       )
             UpperCamelCase    :  Union[str, Any]							=	pa.StructArray.from_arrays([storage, path_array],			['bytes', 'path'],			mask=storage.is_null()       )
          elif pa.types.is_struct(storage.type       ):
             if storage.type.get_field_index('bytes'       ) >= 0:
                UpperCamelCase    :  List[str]							=	storage.field('bytes'       )
             else:
                UpperCamelCase    :  List[str]							=	pa.array([None] * len(SCREAMING_SNAKE_CASE_       ),			type=pa.binary()       )
             if storage.type.get_field_index('path'       ) >= 0:
                UpperCamelCase    :  List[str]							=	storage.field('path'       )
             else:
                UpperCamelCase    :  Optional[Any]							=	pa.array([None] * len(SCREAMING_SNAKE_CASE_       ),			type=pa.string()       )
             UpperCamelCase    :  Optional[int]							=	pa.StructArray.from_arrays([bytes_array, path_array],			['bytes', 'path'],			mask=storage.is_null()       )
          elif pa.types.is_list(storage.type       ):
             UpperCamelCase    :  Optional[Any]							=	pa.array(
                 [encode_np_array(np.array(SCREAMING_SNAKE_CASE_       )       )['bytes'] if arr is not None else None for arr in storage.to_pylist()],			type=pa.binary(),			)
             UpperCamelCase    :  List[str]							=	pa.array([None] * len(SCREAMING_SNAKE_CASE_       ),			type=pa.string()       )
             UpperCamelCase    :  int							=	pa.StructArray.from_arrays(
                 [bytes_array, path_array],			['bytes', 'path'],			mask=bytes_array.is_null()       )
          return array_cast(SCREAMING_SNAKE_CASE_,			self.pa_type       )






       def 					snake_case_       (			self,			SCREAMING_SNAKE_CASE_       )     ->		pa.StructArray:

          @no_op_if_value_is_null
          def path_to_bytes(SCREAMING_SNAKE_CASE_       ):
             with xopen(SCREAMING_SNAKE_CASE_,			'rb'       ) as f:
                UpperCamelCase    :  Optional[int]							=	f.read()
             return bytes_

          UpperCamelCase    :  Union[str, Any]							=	pa.array(
              [
                  (path_to_bytes(x['path']       ) if x['bytes'] is None else x['bytes']) if x is not None else None
                  for x in storage.to_pylist()
              ],			type=pa.binary(),			)
          UpperCamelCase    :  Any							=	pa.array(
              [os.path.basename(SCREAMING_SNAKE_CASE_       ) if path is not None else None for path in storage.field('path'       ).to_pylist()],			type=pa.string(),			)
          UpperCamelCase    :  int							=	pa.StructArray.from_arrays([bytes_array, path_array],			['bytes', 'path'],			mask=bytes_array.is_null()       )
          return array_cast(SCREAMING_SNAKE_CASE_,			self.pa_type       )






def 		UpperCamelCase	(   )	->     List[str]:
   if config.PIL_AVAILABLE:
      import PIL.Image
   else:
      raise ImportError('To support encoding images, please install \'Pillow\'.'	)

   global _IMAGE_COMPRESSION_FORMATS
   if _IMAGE_COMPRESSION_FORMATS is None:
      PIL.Image.init()
      UpperCamelCase    :  Dict							=	list(set(PIL.Image.OPEN.keys()	) & set(PIL.Image.SAVE.keys()	)	)
   return _IMAGE_COMPRESSION_FORMATS
def 		UpperCamelCase	(   snake_case__       :    "PIL.Image.Image"	)	->     bytes:
   UpperCamelCase    :  Any							=	BytesIO()
   if image.format in list_image_compression_formats():
      UpperCamelCase    :  Tuple							=	image.format
   else:
      UpperCamelCase    :  List[str]							=	'PNG' if image.mode in ['1', 'L', 'LA', 'RGB', 'RGBA'] else 'TIFF'
   image.save(snake_case__       ,     format=snake_case__	)
   return buffer.getvalue()
def 		UpperCamelCase	(   snake_case__       :    "PIL.Image.Image"	)	->     dict:
   if hasattr(snake_case__       ,     'filename'	) and image.filename != "":
      return {"path": image.filename, "bytes": None}
   else:
      return {"path": None, "bytes": image_to_bytes(snake_case__	)}
def 		UpperCamelCase	(   snake_case__       :    np.ndarray	)	->     dict:
   if config.PIL_AVAILABLE:
      import PIL.Image
   else:
      raise ImportError('To support encoding images, please install \'Pillow\'.'	)

   UpperCamelCase    :  Union[str, Any]							=	array.dtype
   UpperCamelCase    :  List[Any]							=	dtype.byteorder if dtype.byteorder != '=' else _NATIVE_BYTEORDER
   UpperCamelCase    :  Optional[Any]							=	dtype.kind
   UpperCamelCase    :  Any							=	dtype.itemsize

   UpperCamelCase    :  int							=	None

   # Multi-channel array case (only np.dtype("|u1") is allowed)
   if array.shape[2:]:
      UpperCamelCase    :  Optional[Any]							=	np.dtype('|u1'	)
      if dtype_kind not in ["u", "i"]:
         raise TypeError(
             F"""Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays."""	)
      if dtype is not dest_dtype:
         warnings.warn(F"""Downcasting array dtype {dtype} to {dest_dtype} to be compatible with 'Pillow'"""	)
    # Exact match
   elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
      UpperCamelCase    :  List[Any]							=	dtype
   else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
      while dtype_itemsize >= 1:
         UpperCamelCase    :  Dict							=	dtype_byteorder + dtype_kind + str(snake_case__	)
         UpperCamelCase    :  str							=	np.dtype(snake_case__	)
         if dest_dtype in _VALID_IMAGE_ARRAY_DTPYES:
            warnings.warn(F"""Downcasting array dtype {dtype} to {dest_dtype} to be compatible with 'Pillow'"""	)
            break
         else:
            dtype_itemsize //= 2

   if dest_dtype is None:
      raise TypeError(
          F"""Cannot convert dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}"""	)
   UpperCamelCase    :  Union[str, Any]							=	PIL.Image.fromarray(array.astype(snake_case__	)	)
   return {"path": None, "bytes": image_to_bytes(snake_case__	)}


def 		UpperCamelCase	(   snake_case__       :    Union[List[str], List[dict], List[np.ndarray], List["PIL.Image.Image"]]	)	->     List[dict]:
   if config.PIL_AVAILABLE:
      import PIL.Image
   else:
      raise ImportError('To support encoding images, please install \'Pillow\'.'	)

   if objs:
      UpperCamelCase	,     UpperCamelCase    :  Union[str, Any]							=	first_non_null_value(snake_case__	)
      if isinstance(snake_case__       ,     snake_case__	):
         return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
      if isinstance(snake_case__       ,     np.ndarray	):
         UpperCamelCase    :  List[Any]							=	no_op_if_value_is_null(snake_case__	)
         return [obj_to_image_dict_func(snake_case__	) for obj in objs]
      elif isinstance(snake_case__       ,     PIL.Image.Image	):
         UpperCamelCase    :  Optional[int]							=	no_op_if_value_is_null(snake_case__	)
         return [obj_to_image_dict_func(snake_case__	) for obj in objs]
      else:
         return objs
   else:
      return objs



