




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


_UpperCAmelCase		:  List[Any]	=   logging.getLogger(__name__)




class       __lowerCAmelCase				(							lowerCAmelCase):






       def __init__(      self:    Dict				,    _lowerCAmelCase:    Any				,    _lowerCAmelCase:    List[Any]				,    _lowerCAmelCase:    Any				,    _lowerCAmelCase:    Any=None ):
         super().__init__(
             _lowerCAmelCase				,    question_encoder_tokenizer=_lowerCAmelCase				,    generator_tokenizer=_lowerCAmelCase				,    index=_lowerCAmelCase				,    init_retrieval=_lowerCAmelCase				,    )
         lowercase			:int							=  None






       def     SCREAMING_SNAKE_CASE			(      self:    Dict				,    _lowerCAmelCase:    int ):
         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
           lowercase			:Tuple							=  self._infer_socket_ifname()
           # avoid clash with the NCCL port
           lowercase			:Union[str, Any]							=  str(distributed_port + 1 )
           lowercase			:Optional[Any]							=  dist.new_group(ranks=_lowerCAmelCase				,    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     SCREAMING_SNAKE_CASE			(      self:    Optional[Any] ):
         return dist.get_rank(group=self.process_group ) == 0






       def     SCREAMING_SNAKE_CASE			(      self:    int				,    _lowerCAmelCase:    List[Any]				,    _lowerCAmelCase:    List[str]				,    _lowerCAmelCase:    Optional[Any]=torch.floataa ):
         lowercase			:Dict							=  torch.empty(_lowerCAmelCase				,    dtype=_lowerCAmelCase )
         dist.scatter(_lowerCAmelCase				,    src=0				,    scatter_list=_lowerCAmelCase				,    group=self.process_group )
         return target_tensor






       def     SCREAMING_SNAKE_CASE			(      self:    List[Any] ):
         lowercase			:int							=  psutil.net_if_addrs()
         # a hacky way to deal with varying network interface names
         lowercase			:str							=  next((addr for addr in addrs if addr.startswith("e" ))				,    _lowerCAmelCase )
         return ifname






       def     SCREAMING_SNAKE_CASE			(      self:    List[str]				,    _lowerCAmelCase:    np.ndarray				,    _lowerCAmelCase:    int ):

         # single GPU training
         if not dist.is_initialized():
           lowercase		,			lowercase			:Any							=  self._main_retrieve(_lowerCAmelCase				,    _lowerCAmelCase )
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_lowerCAmelCase )

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

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

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

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









from typing import TYPE_CHECKING

from ....utils import _LazyModule


_UpperCAmelCase		:  Dict	=   {"tokenization_tapex": ["TapexTokenizer"]}


if TYPE_CHECKING:
     from .tokenization_tapex import TapexTokenizer


else:
     import sys

     _UpperCAmelCase		:  Optional[Any]	=   _LazyModule(__name__, globals()["__file__"], _import_structure)



