


'''simple docstring'''






import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





def lowerCamelCase__      (						__lowerCamelCase	:							Dict	,   __lowerCamelCase	:							int							):







   '''simple docstring'''




   assert isinstance(__lowerCamelCase	,   __lowerCamelCase							)
   assert dataset.num_rows == 4
   assert dataset.num_columns == 3
   assert dataset.column_names == ["col_1", "col_2", "col_3"]
   for feature, expected_dtype in expected_features.items():
      assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize('keep_in_memory'	,   [False, True]							)
def lowerCamelCase__      (						__lowerCamelCase	:							Optional[Any]	,   __lowerCamelCase	:							Optional[int]	,   __lowerCamelCase	:							Optional[Any]							):







   '''simple docstring'''




   _UpperCAmelCase       :       Optional[Any]							   =tmp_path / 'cache'
   _UpperCAmelCase       :       Optional[Any]							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      _UpperCAmelCase       :       Optional[Any]							   =ParquetDatasetReader(__lowerCamelCase	,   cache_dir=__lowerCamelCase	,   keep_in_memory=__lowerCamelCase							).read()
   _check_parquet_dataset(__lowerCamelCase	,   __lowerCamelCase							)





@pytest.mark.parametrize(
    'features'	,   [
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ]	,   )
def lowerCamelCase__      (						__lowerCamelCase	:							str	,   __lowerCamelCase	:							List[str]	,   __lowerCamelCase	:							Any							):







   '''simple docstring'''




   _UpperCAmelCase       :       Optional[Any]							   =tmp_path / 'cache'
   _UpperCAmelCase       :       List[Any]							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   _UpperCAmelCase       :       Optional[Any]							   =features.copy() if features else default_expected_features
   _UpperCAmelCase       :       List[str]							   =(
       Features({feature: Value(__lowerCamelCase							) for feature, dtype in features.items()}							) if features is not None else None
   )
   _UpperCAmelCase       :       int							   =ParquetDatasetReader(__lowerCamelCase	,   features=__lowerCamelCase	,   cache_dir=__lowerCamelCase							).read()
   _check_parquet_dataset(__lowerCamelCase	,   __lowerCamelCase							)





@pytest.mark.parametrize('split'	,   [None, NamedSplit('train'							), 'train', 'test']							)
def lowerCamelCase__      (						__lowerCamelCase	:							str	,   __lowerCamelCase	:							Optional[Any]	,   __lowerCamelCase	:							Optional[Any]							):







   '''simple docstring'''




   _UpperCAmelCase       :       Any							   =tmp_path / 'cache'
   _UpperCAmelCase       :       int							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   _UpperCAmelCase       :       List[str]							   =ParquetDatasetReader(__lowerCamelCase	,   cache_dir=__lowerCamelCase	,   split=__lowerCamelCase							).read()
   _check_parquet_dataset(__lowerCamelCase	,   __lowerCamelCase							)
   assert dataset.split == split if split else "train"





@pytest.mark.parametrize('path_type'	,   [str, list]							)
def lowerCamelCase__      (						__lowerCamelCase	:							Optional[int]	,   __lowerCamelCase	:							Dict	,   __lowerCamelCase	:							Union[str, Any]							):







   '''simple docstring'''




   if issubclass(__lowerCamelCase	,   __lowerCamelCase							):
      _UpperCAmelCase       :       Union[str, Any]							   =parquet_path
   elif issubclass(__lowerCamelCase	,   __lowerCamelCase							):
      _UpperCAmelCase       :       Union[str, Any]							   =[parquet_path]
   _UpperCAmelCase       :       Union[str, Any]							   =tmp_path / 'cache'
   _UpperCAmelCase       :       Union[str, Any]							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   _UpperCAmelCase       :       int							   =ParquetDatasetReader(__lowerCamelCase	,   cache_dir=__lowerCamelCase							).read()
   _check_parquet_dataset(__lowerCamelCase	,   __lowerCamelCase							)





def lowerCamelCase__      (						__lowerCamelCase	:							List[str]	,   __lowerCamelCase	:							Union[str, Any]	,   __lowerCamelCase	:							str=("train",)							):







   '''simple docstring'''




   assert isinstance(__lowerCamelCase	,   __lowerCamelCase							)
   for split in splits:
      _UpperCAmelCase       :       List[Any]							   =dataset_dict[split]
      assert dataset.num_rows == 4
      assert dataset.num_columns == 3
      assert dataset.column_names == ["col_1", "col_2", "col_3"]
      for feature, expected_dtype in expected_features.items():
         assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize('keep_in_memory'	,   [False, True]							)
def lowerCamelCase__      (						__lowerCamelCase	:							Union[str, Any]	,   __lowerCamelCase	:							int	,   __lowerCamelCase	:							str							):







   '''simple docstring'''




   _UpperCAmelCase       :       Any							   =tmp_path / 'cache'
   _UpperCAmelCase       :       Optional[int]							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      _UpperCAmelCase       :       List[Any]							   =ParquetDatasetReader(
          {'train': parquet_path}	,   cache_dir=__lowerCamelCase	,   keep_in_memory=__lowerCamelCase							).read()
   _check_parquet_datasetdict(__lowerCamelCase	,   __lowerCamelCase							)





@pytest.mark.parametrize(
    'features'	,   [
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ]	,   )
def lowerCamelCase__      (						__lowerCamelCase	:							Optional[int]	,   __lowerCamelCase	:							List[str]	,   __lowerCamelCase	:							Union[str, Any]							):







   '''simple docstring'''




   _UpperCAmelCase       :       str							   =tmp_path / 'cache'
   _UpperCAmelCase       :       Tuple							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   _UpperCAmelCase       :       Any							   =features.copy() if features else default_expected_features
   _UpperCAmelCase       :       Optional[Any]							   =(
       Features({feature: Value(__lowerCamelCase							) for feature, dtype in features.items()}							) if features is not None else None
   )
   _UpperCAmelCase       :       Tuple							   =ParquetDatasetReader({'train': parquet_path}	,   features=__lowerCamelCase	,   cache_dir=__lowerCamelCase							).read()
   _check_parquet_datasetdict(__lowerCamelCase	,   __lowerCamelCase							)





@pytest.mark.parametrize('split'	,   [None, NamedSplit('train'							), 'train', 'test']							)
def lowerCamelCase__      (						__lowerCamelCase	:							Tuple	,   __lowerCamelCase	:							Union[str, Any]	,   __lowerCamelCase	:							Optional[int]							):







   '''simple docstring'''




   if split:
      _UpperCAmelCase       :       Dict							   ={split: parquet_path}
   else:
      _UpperCAmelCase       :       Optional[Any]							   ='train'
      _UpperCAmelCase       :       Any							   ={'train': parquet_path, 'test': parquet_path}
   _UpperCAmelCase       :       Any							   =tmp_path / 'cache'
   _UpperCAmelCase       :       Union[str, Any]							   ={'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   _UpperCAmelCase       :       List[Any]							   =ParquetDatasetReader(__lowerCamelCase	,   cache_dir=__lowerCamelCase							).read()
   _check_parquet_datasetdict(__lowerCamelCase	,   __lowerCamelCase	,   splits=list(path.keys()							)							)
   assert all(dataset[split].split == split for split in path.keys()							)





def lowerCamelCase__      (						__lowerCamelCase	:							Optional[int]	,   __lowerCamelCase	:							str							):







   '''simple docstring'''




   _UpperCAmelCase       :       str							   =ParquetDatasetWriter(__lowerCamelCase	,   tmp_path / 'foo.parquet'							)
   assert writer.write() > 0
   _UpperCAmelCase       :       Dict							   =pq.ParquetFile(tmp_path / 'foo.parquet'							)
   _UpperCAmelCase       :       str							   =pf.read()
   assert dataset.data.table == output_table





def lowerCamelCase__      (						__lowerCamelCase	:							Optional[Any]	,   __lowerCamelCase	:							Any							):







   '''simple docstring'''




   _UpperCAmelCase       :       List[str]							   =str(shared_datadir / 'test_image_rgb.jpg'							)
   _UpperCAmelCase       :       Any							   ={'image': [image_path]}
   _UpperCAmelCase       :       int							   =Features({'image': Image()}							)
   _UpperCAmelCase       :       Any							   =Dataset.from_dict(__lowerCamelCase	,   features=__lowerCamelCase							)
   _UpperCAmelCase       :       Dict							   =ParquetDatasetWriter(__lowerCamelCase	,   tmp_path / 'foo.parquet'							)
   assert writer.write() > 0

   _UpperCAmelCase       :       Any							   =Dataset.from_parquet(str(tmp_path / 'foo.parquet'							)							)
   assert dataset.features == reloaded_dataset.features

   _UpperCAmelCase       :       Dict							   =ParquetDatasetReader(str(tmp_path / 'foo.parquet'							)	,   streaming=__lowerCamelCase							).read()
   assert dataset.features == reloaded_iterable_dataset.features





@pytest.mark.parametrize(
    'feature, expected'	,   [
        (Features({'foo': Value('int32'							)}							), None),
        (Features({'image': Image(), 'foo': Value('int32'							)}							), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({'nested': Sequence(Audio()							)}							), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ]	,   )
def lowerCamelCase__      (						__lowerCamelCase	:							List[Any]	,   __lowerCamelCase	:							List[Any]							):







   '''simple docstring'''




   assert get_writer_batch_size(__lowerCamelCase							) == expected





'''simple docstring'''






import logging
import os
from typing import List, Tuple

import numpy as np
import psutil
import torch
import torch.distributed as dist

from transformers import RagRetriever


lowercase							=logging.getLogger(__name__)






class  __magic_name__							(		lowerCAmelCase    ):






       def __init__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case=None)						->  List[Any]:




          '''simple docstring'''
          super().__init__(
              snake_case  ,							question_encoder_tokenizer=snake_case  ,							generator_tokenizer=snake_case  ,							index=snake_case  ,							init_retrieval=snake_case  ,							)
          _UpperCAmelCase       :       Union[str, Any]							   =None






       def       lowerCAmelCase		(      self  ,							snake_case)						->  Optional[Any]:




          '''simple docstring'''
          logger.info('initializing retrieval')

          # initializing a separate process group for retrieval as the default
          # nccl backend doesn't support gather/scatter operations while gloo
          # is too slow to replace nccl for the core gpu communication
          if dist.is_initialized():
             logger.info('dist initialized')
             # needs to be set manually
             _UpperCAmelCase       :       Optional[Any]							   =self._infer_socket_ifname()
             # avoid clash with the NCCL port
             _UpperCAmelCase       :       Optional[int]							   =str(distributed_port + 1)
             _UpperCAmelCase       :       Any							   =dist.new_group(ranks=snake_case  ,							backend='gloo')

          # initialize retriever only on the main worker
          if not dist.is_initialized() or self._is_main():
             logger.info('dist not initialized / main')
             self.index.init_index()

          # all processes wait untill the retriever is initialized by the main process
          if dist.is_initialized():
             torch.distributed.barrier(group=self.process_group)






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          return dist.get_rank(group=self.process_group) == 0






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case  ,							snake_case=torch.floataa)						->  Union[str, Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =torch.empty(snake_case  ,							dtype=snake_case)
          dist.scatter(snake_case  ,							src=0  ,							scatter_list=snake_case  ,							group=self.process_group)
          return target_tensor






       def       lowerCAmelCase		(      self)						->  Optional[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       Dict							   =psutil.net_if_addrs()
          # a hacky way to deal with varying network interface names
          _UpperCAmelCase       :       str							   =next((addr for addr in addrs if addr.startswith('e'))  ,							snake_case)
          return ifname






       def       lowerCAmelCase		(      self  ,							snake_case  ,							snake_case)						->  Tuple[np.ndarray, List[dict]]:




          '''simple docstring'''

          # single GPU training
          if not dist.is_initialized():
             _UpperCAmelCase   ,	_UpperCAmelCase       :       Optional[Any]							   =self._main_retrieve(snake_case  ,							snake_case)
             return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(snake_case)

          # distributed training
          _UpperCAmelCase       :       Optional[int]							   =dist.get_world_size(group=self.process_group)

          # gather logic
          _UpperCAmelCase       :       str							   =None
          if self._is_main():
             _UpperCAmelCase       :       Union[str, Any]							   =[torch.empty(question_hidden_states.shape  ,							dtype=torch.floataa) for _ in range(snake_case)]
          dist.gather(torch.tensor(snake_case)  ,							dst=0  ,							gather_list=snake_case  ,							group=self.process_group)

          # scatter logic
          _UpperCAmelCase       :       Optional[Any]							   =question_hidden_states.shape[0]
          _UpperCAmelCase       :       List[Any]							   =[]
          _UpperCAmelCase       :       Any							   =[]
          if self._is_main():
             assert len(snake_case) == world_size
             _UpperCAmelCase   ,	_UpperCAmelCase       :       Tuple							   =self._main_retrieve(torch.cat(snake_case).numpy()  ,							snake_case)
             _UpperCAmelCase   ,	_UpperCAmelCase       :       Union[str, Any]							   =torch.tensor(snake_case), torch.tensor(snake_case)
             _UpperCAmelCase       :       List[str]							   =self._chunk_tensor(snake_case  ,							snake_case)
             _UpperCAmelCase       :       Union[str, Any]							   =self._chunk_tensor(snake_case  ,							snake_case)
          _UpperCAmelCase       :       int							   =self._scattered(snake_case  ,							[n_queries, n_docs]  ,							target_type=torch.intaa)
          _UpperCAmelCase       :       Dict							   =self._scattered(snake_case  ,							[n_queries, n_docs, question_hidden_states.shape[1]])

          return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(snake_case)

