






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


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

class     __SCREAMING_SNAKE_CASE					(   __lowercase):






						def __init__(	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase=None   ):






												"""simple docstring"""
												super().__init__(
												    _UpperCamelCase     ,						question_encoder_tokenizer=_UpperCamelCase     ,						generator_tokenizer=_UpperCamelCase     ,						index=_UpperCamelCase     ,						init_retrieval=_UpperCamelCase     ,						)
												lowerCAmelCase__    						=						None






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase   ):






												"""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__    						=						self._infer_socket_ifname()
																		# avoid clash with the NCCL port
																		lowerCAmelCase__    						=						str(distributed_port + 1   )
																		lowerCAmelCase__    						=						dist.new_group(ranks=_UpperCamelCase     ,						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       UpperCamelCase__      (	self   ):






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






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase     ,						_UpperCamelCase=torch.floataa   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						torch.empty(_UpperCamelCase     ,						dtype=_UpperCamelCase   )
												dist.scatter(_UpperCamelCase     ,						src=0     ,						scatter_list=_UpperCamelCase     ,						group=self.process_group   )
												return target_tensor






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						psutil.net_if_addrs()
												# a hacky way to deal with varying network interface names
												lowerCAmelCase__    						=						next((addr for addr in addrs if addr.startswith('e'   ))     ,						_UpperCamelCase   )
												return ifname







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""

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

												# distributed training
												lowerCAmelCase__    						=						dist.get_world_size(group=self.process_group   )

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

												# scatter logic
												lowerCAmelCase__    						=						question_hidden_states.shape[0]
												lowerCAmelCase__    						=						[]
												lowerCAmelCase__    						=						[]
												if self._is_main():
																		assert len(_UpperCamelCase   ) == world_size
																		lowerCAmelCase__						,       lowerCAmelCase__    						=						self._main_retrieve(torch.cat(_UpperCamelCase   ).numpy()     ,						_UpperCamelCase   )
																		lowerCAmelCase__						,       lowerCAmelCase__    						=						torch.tensor(_UpperCamelCase   ), torch.tensor(_UpperCamelCase   )
																		lowerCAmelCase__    						=						self._chunk_tensor(_UpperCamelCase     ,						_UpperCamelCase   )
																		lowerCAmelCase__    						=						self._chunk_tensor(_UpperCamelCase     ,						_UpperCamelCase   )
												lowerCAmelCase__    						=						self._scattered(_UpperCamelCase     ,						[n_queries, n_docs]     ,						target_type=torch.intaa   )
												lowerCAmelCase__    						=						self._scattered(_UpperCamelCase     ,						[n_queries, n_docs, question_hidden_states.shape[1]]   )

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













import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
					from transformers.pipelines.conversational import Conversation


__snake_case    :    Any         =     logging.get_logger(__name__)

__snake_case    :    Tuple         =     {"""tokenizer_file""": """tokenizer.json"""}

__snake_case    :    str         =     {
    """tokenizer_file""": {
        """bigscience/tokenizer""": """https://huggingface.co/bigscience/tokenizer/blob/main/tokenizer.json""",
        """bigscience/bloom-560m""": """https://huggingface.co/bigscience/bloom-560m/blob/main/tokenizer.json""",
        """bigscience/bloom-1b1""": """https://huggingface.co/bigscience/bloom-1b1/blob/main/tokenizer.json""",
        """bigscience/bloom-1b7""": """https://huggingface.co/bigscience/bloom-1b7/blob/main/tokenizer.json""",
        """bigscience/bloom-3b""": """https://huggingface.co/bigscience/bloom-3b/blob/main/tokenizer.json""",
        """bigscience/bloom-7b1""": """https://huggingface.co/bigscience/bloom-7b1/blob/main/tokenizer.json""",
        """bigscience/bloom""": """https://huggingface.co/bigscience/bloom/blob/main/tokenizer.json""",
    },
}


class     __SCREAMING_SNAKE_CASE					(   __lowercase):
						_SCREAMING_SNAKE_CASE	:       Dict   =					VOCAB_FILES_NAMES
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					PRETRAINED_VOCAB_FILES_MAP
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					['''input_ids''', '''attention_mask''']
						_SCREAMING_SNAKE_CASE	:       Union[str, Any]   =					None






						def __init__(	self     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase=None     ,						_UpperCamelCase="<unk>"     ,						_UpperCamelCase="<s>"     ,						_UpperCamelCase="</s>"     ,						_UpperCamelCase="<pad>"     ,						_UpperCamelCase=False     ,						_UpperCamelCase=False     ,						**_UpperCamelCase     ,						):






												"""simple docstring"""
												super().__init__(
												    _UpperCamelCase     ,						_UpperCamelCase     ,						tokenizer_file=_UpperCamelCase     ,						unk_token=_UpperCamelCase     ,						bos_token=_UpperCamelCase     ,						eos_token=_UpperCamelCase     ,						pad_token=_UpperCamelCase     ,						add_prefix_space=_UpperCamelCase     ,						clean_up_tokenization_spaces=_UpperCamelCase     ,						**_UpperCamelCase     ,						)
												lowerCAmelCase__    						=						json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()   )
												if pre_tok_state.get('add_prefix_space'     ,						_UpperCamelCase   ) != add_prefix_space:
																		lowerCAmelCase__    						=						getattr(_UpperCamelCase     ,						pre_tok_state.pop('type'   )   )
																		lowerCAmelCase__    						=						add_prefix_space
																		lowerCAmelCase__    						=						pre_tok_class(**_UpperCamelCase   )

												lowerCAmelCase__    						=						add_prefix_space






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )
												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._batch_encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						kwargs.get('is_split_into_words'     ,						_UpperCamelCase   )

												if not (self.add_prefix_space or not is_split_into_words):
																		raise Exception(
																		    F"You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"
																		    ' pretokenized inputs.'   )

												return super()._encode_plus(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase = None   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						self._tokenizer.model.save(_UpperCamelCase     ,						name=_UpperCamelCase   )
												return tuple(_UpperCamelCase   )







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						[]
												for is_user, text in conversation.iter_texts():
																		input_ids.extend(self.encode(_UpperCamelCase     ,						add_special_tokens=_UpperCamelCase   ) + [self.eos_token_id]   )

												if len(_UpperCamelCase   ) > self.model_max_length:
																		lowerCAmelCase__    						=						input_ids[-self.model_max_length :]
												return input_ids





