



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


lowerCAmelCase__      =    logging.getLogger(__name__)






class 	__snake_case	(     _lowercase):




  def __init__(			self					:						Union[str, Any]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						int=None		):






         """simple docstring"""






         super().__init__(
             __lowerCAmelCase						, question_encoder_tokenizer=__lowerCAmelCase						, generator_tokenizer=__lowerCAmelCase						, index=__lowerCAmelCase						, init_retrieval=__lowerCAmelCase						, )
         _lowerCamelCase						:    Dict         =			None




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict						, __lowerCAmelCase					:						int		):






         """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
                _lowerCamelCase						:    List[str]         =			self._infer_socket_ifname()
                # avoid clash with the NCCL port
                _lowerCamelCase						:    Dict         =			str(distributed_port + 1		)
                _lowerCamelCase						:    str         =			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]		):






         """simple docstring"""






         return dist.get_rank(group=self.process_group		) == 0




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Tuple=torch.floataa		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			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					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			psutil.net_if_addrs()
         # a hacky way to deal with varying network interface names
         _lowerCamelCase						:    str         =			next((addr for addr in addrs if addr.startswith('''e'''		))						, __lowerCAmelCase		)
         return ifname




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						np.ndarray						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         if not dist.is_initialized():
                _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			self._main_retrieve(__lowerCAmelCase						, __lowerCAmelCase		)
                return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__lowerCAmelCase		)

         # distributed training
         _lowerCamelCase						:    Dict         =			dist.get_world_size(group=self.process_group		)

         # gather logic
         _lowerCamelCase						:    str         =			None
         if self._is_main():
                _lowerCamelCase						:    List[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
         _lowerCamelCase						:    int         =			question_hidden_states.shape[0]
         _lowerCamelCase						:    str         =			[]
         _lowerCamelCase						:    Optional[int]         =			[]
         if self._is_main():
                assert len(__lowerCAmelCase		) == world_size
                _lowerCamelCase      ,   _lowerCamelCase						:    Tuple         =			self._main_retrieve(torch.cat(__lowerCAmelCase		).numpy()						, __lowerCAmelCase		)
                _lowerCamelCase      ,   _lowerCamelCase						:    List[str]         =			torch.tensor(__lowerCAmelCase		), torch.tensor(__lowerCAmelCase		)
                _lowerCamelCase						:    List[Any]         =			self._chunk_tensor(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    Optional[Any]         =			self._chunk_tensor(__lowerCAmelCase						, __lowerCAmelCase		)
         _lowerCamelCase						:    int         =			self._scattered(__lowerCAmelCase						, [n_queries, n_docs]						, target_type=torch.intaa		)
         _lowerCamelCase						:    str         =			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		)







"""simple docstring"""






import json
import os
import unittest

from transformers.models.blenderbot_small.tokenization_blenderbot_small import (
    VOCAB_FILES_NAMES,
    BlenderbotSmallTokenizer,
)

from ...test_tokenization_common import TokenizerTesterMixin






class 	__snake_case	(     _lowercase  ,					unittest.TestCase):
  snake_case__    :       Optional[int]        =						BlenderbotSmallTokenizer
  snake_case__    :       List[str]        =						False




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         super().setUp()

         _lowerCamelCase						:    str         =			['''__start__''', '''adapt''', '''act''', '''ap@@''', '''te''', '''__end__''', '''__unk__''']
         _lowerCamelCase						:    Any         =			dict(zip(__lowerCAmelCase						, range(len(__lowerCAmelCase		)		)		)		)

         _lowerCamelCase						:    Any         =			['''#version: 0.2''', '''a p''', '''t e</w>''', '''ap t</w>''', '''a d''', '''ad apt</w>''', '''a c''', '''ac t</w>''', '''''']
         _lowerCamelCase						:    List[str]         =			{'''unk_token''': '''__unk__''', '''bos_token''': '''__start__''', '''eos_token''': '''__end__'''}

         _lowerCamelCase						:    Dict         =			os.path.join(self.tmpdirname						, VOCAB_FILES_NAMES['''vocab_file''']		)
         _lowerCamelCase						:    List[str]         =			os.path.join(self.tmpdirname						, VOCAB_FILES_NAMES['''merges_file''']		)
         with open(self.vocab_file						, '''w'''						, encoding='''utf-8'''		) as fp:
                fp.write(json.dumps(__lowerCAmelCase		) + '''\n'''		)
         with open(self.merges_file						, '''w'''						, encoding='''utf-8'''		) as fp:
                fp.write('''\n'''.join(__lowerCAmelCase		)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, **__lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         kwargs.update(self.special_tokens_map		)
         return BlenderbotSmallTokenizer.from_pretrained(self.tmpdirname						, **__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			'''adapt act apte'''
         _lowerCamelCase						:    Tuple         =			'''adapt act apte'''
         return input_text, output_text




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			BlenderbotSmallTokenizer(self.vocab_file						, self.merges_file						, **self.special_tokens_map		)
         _lowerCamelCase						:    int         =			'''adapt act apte'''
         _lowerCamelCase						:    Optional[Any]         =			['''adapt''', '''act''', '''ap@@''', '''te''']
         _lowerCamelCase						:    Optional[Any]         =			tokenizer.tokenize(__lowerCAmelCase		)
         self.assertListEqual(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			[tokenizer.bos_token] + tokens + [tokenizer.eos_token]

         _lowerCamelCase						:    Dict         =			[0, 1, 2, 3, 4, 5]
         self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCAmelCase		)						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    str         =			BlenderbotSmallTokenizer.from_pretrained('''facebook/blenderbot-90M'''		)
         assert tok('''sam'''		).input_ids == [1_3_8_4]
         _lowerCamelCase						:    List[str]         =			'''I am a small frog.'''
         _lowerCamelCase						:    str         =			tok([src_text]						, padding=__lowerCAmelCase						, truncation=__lowerCAmelCase		)['''input_ids''']
         _lowerCamelCase						:    Any         =			tok.batch_decode(__lowerCAmelCase						, skip_special_tokens=__lowerCAmelCase						, clean_up_tokenization_spaces=__lowerCAmelCase		)[0]
         assert src_text != decoded  # I wish it did!
         assert decoded == "i am a small frog ."




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			BlenderbotSmallTokenizer.from_pretrained('''facebook/blenderbot-90M'''		)
         _lowerCamelCase						:    Optional[Any]         =			'''I am a small frog .'''
         _lowerCamelCase						:    str         =			'''.'''
         _lowerCamelCase						:    str         =			tok(__lowerCAmelCase		)['''input_ids''']
         _lowerCamelCase						:    Dict         =			tok(__lowerCAmelCase		)['''input_ids''']

         assert encoded[-1] == encoded_dot[0]


