





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


_lowerCAmelCase 				=     logging.getLogger(__name__)

class    _UpperCAmelCase							(		_lowerCamelCase     ):

      '''simple docstring'''





      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__=None       ):
         super().__init__(
             __A      ,       question_encoder_tokenizer=__A      ,       generator_tokenizer=__A      ,       index=__A      ,       init_retrieval=__A      ,       )
         A_     :				Union[str, Any]			      =	None
      def 		_lowerCamelCase      (		self      ,       a__       ):
         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
            A_     :				List[str]			      =	self._infer_socket_ifname()
            # avoid clash with the NCCL port
            A_     :				int			      =	str(distributed_port + 1       )
            A_     :				Any			      =	dist.new_group(ranks=__A      ,       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       ):
         return dist.get_rank(group=self.process_group       ) == 0
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__=torch.floataa       ):
         A_     :				Dict			      =	torch.empty(__A      ,       dtype=__A       )
         dist.scatter(__A      ,       src=0      ,       scatter_list=__A      ,       group=self.process_group       )
         return target_tensor
      def 		_lowerCamelCase      (		self       ):
         A_     :				int			      =	psutil.net_if_addrs()
         # a hacky way to deal with varying network interface names
         A_     :				Optional[int]			      =	next((addr for addr in addrs if addr.startswith("""e"""       ))      ,       __A       )
         return ifname






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__       ):

         # single GPU training
         if not dist.is_initialized():
            A_     :				Dict			      =	self._main_retrieve(__A      ,       __A       )
            return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__A       )

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

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

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

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













from typing import List

from .keymap import KEYMAP, get_character



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
   '''simple docstring'''


   def decorator(_lowerCAmelCase     ):
      A_     :				List[Any]			      =	getattr(_lowerCAmelCase  ,"""handle_key"""  ,[]     )
      handle += [key]
      setattr(_lowerCAmelCase  ,"""handle_key"""  ,_lowerCAmelCase     )
      return func

   return decorator



def     _lowerCAmelCase		(  *_lowerCAmelCase     ):
   '''simple docstring'''


   def decorator(_lowerCAmelCase     ):
      A_     :				Tuple			      =	getattr(_lowerCAmelCase  ,"""handle_key"""  ,[]     )
      handle += keys
      setattr(_lowerCAmelCase  ,"""handle_key"""  ,_lowerCAmelCase     )
      return func

   return decorator

class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __new__(		cls      ,       a__      ,       a__      ,       a__       ):
         A_     :				Any			      =	super().__new__(cls      ,       a__      ,       a__      ,       a__       )
         if not hasattr(a__      ,       """key_handler"""       ):
            setattr(a__      ,       """key_handler"""      ,       {}       )
         setattr(a__      ,       """handle_input"""      ,       KeyHandler.handle_input       )

         for value in attrs.values():
            A_     :				Tuple			      =	getattr(a__      ,       """handle_key"""      ,       []       )
            for key in handled_keys:
               A_     :				Optional[Any]			      =	value
         return new_cls






      @staticmethod
      def 		_lowerCamelCase      (		cls       ):
         A_     :				List[str]			      =	get_character()
         if char != KEYMAP["undefined"]:
            A_     :				str			      =	ord(a__       )
         A_     :				List[str]			      =	cls.key_handler.get(a__       )
         if handler:
            A_     :				Optional[int]			      =	char
            return handler(cls       )
         else:
            return None






def     _lowerCAmelCase		(  cls     ):
   '''simple docstring'''


   return KeyHandler(cls.__name__  ,cls.__bases__  ,cls.__dict__.copy()     )






