


'''simple docstring'''






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

						if is_vision_available():
												from transformers import MaskaFormerImageProcessor

if is_vision_available():
						from PIL import Image




class 	a__       :





							'''simple docstring'''






							def __init__(	self      ,					lowerCamelCase_      ,					lowerCamelCase_=2      ,					lowerCamelCase_=True      ,					lowerCamelCase_=False      ,					lowerCamelCase_=10      ,					lowerCamelCase_=3      ,					lowerCamelCase_=32 * 8      ,					lowerCamelCase_=32 * 8      ,					lowerCamelCase_=4      ,					lowerCamelCase_=64      ,					) ->							List[Any]:
														lowerCAmelCase__      =			parent
														lowerCAmelCase__      =			batch_size
														lowerCAmelCase__      =			is_training
														lowerCAmelCase__      =			use_auxiliary_loss
														lowerCAmelCase__      =			num_queries
														lowerCAmelCase__      =			num_channels
														lowerCAmelCase__      =			min_size
														lowerCAmelCase__      =			max_size
														lowerCAmelCase__      =			num_labels
														lowerCAmelCase__      =			hidden_dim
														lowerCAmelCase__      =			hidden_dim







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							int:
														lowerCAmelCase__      =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
														    _SCREAMING_SNAKE_CASE     )

														lowerCAmelCase__      =			torch.ones([self.batch_size, self.min_size, self.max_size]      ,					device=_SCREAMING_SNAKE_CASE     )

														lowerCAmelCase__      =			(
														    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]      ,					device=_SCREAMING_SNAKE_CASE     ) > 0.5
														).float()
														lowerCAmelCase__      =			(torch.rand((self.batch_size, self.num_labels)      ,					device=_SCREAMING_SNAKE_CASE     ) > 0.5).long()

														lowerCAmelCase__      =			self.get_config()
														return config, pixel_values, pixel_mask, mask_labels, class_labels







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			MaskaFormerConfig(
														    hidden_size=self.hidden_dim      ,					)
														lowerCAmelCase__      =			self.num_queries
														lowerCAmelCase__      =			self.num_labels

														lowerCAmelCase__      =			[1, 1, 1, 1]
														lowerCAmelCase__      =			self.num_channels

														lowerCAmelCase__      =			64
														lowerCAmelCase__      =			1_28
														lowerCAmelCase__      =			self.hidden_dim
														lowerCAmelCase__      =			self.hidden_dim
														lowerCAmelCase__      =			self.hidden_dim
														return config







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							str:
														lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.prepare_config_and_inputs()
														lowerCAmelCase__      =			{'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
														return config, inputs_dict







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							Union[str, Any]:
														lowerCAmelCase__      =			output.encoder_hidden_states
														lowerCAmelCase__      =			output.pixel_decoder_hidden_states
														lowerCAmelCase__      =			output.transformer_decoder_hidden_states

														self.parent.assertTrue(len(_SCREAMING_SNAKE_CASE     )      ,					len(config.backbone_config.depths     )     )
														self.parent.assertTrue(len(_SCREAMING_SNAKE_CASE     )      ,					len(config.backbone_config.depths     )     )
														self.parent.assertTrue(len(_SCREAMING_SNAKE_CASE     )      ,					config.decoder_layers     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_=False     ) ->							Tuple:






														with torch.no_grad():
																					lowerCAmelCase__      =			MaskaFormerModel(config=_SCREAMING_SNAKE_CASE     )
																					model.to(_SCREAMING_SNAKE_CASE     )
																					model.eval()

																					lowerCAmelCase__      =			model(pixel_values=_SCREAMING_SNAKE_CASE      ,					pixel_mask=_SCREAMING_SNAKE_CASE     )
																					lowerCAmelCase__      =			model(_SCREAMING_SNAKE_CASE      ,					output_hidden_states=_SCREAMING_SNAKE_CASE     )

														self.parent.assertEqual(
														    output.transformer_decoder_last_hidden_state.shape      ,					(self.batch_size, self.num_queries, self.hidden_dim)      ,					)
														# let's ensure the other two hidden state exists
														self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None     )
														self.parent.assertTrue(output.encoder_last_hidden_state is not None     )

														if output_hidden_states:
																					self.check_output_hidden_state(_SCREAMING_SNAKE_CASE      ,					_SCREAMING_SNAKE_CASE     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							List[Any]:
														lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation(config=_SCREAMING_SNAKE_CASE     )
														model.to(_SCREAMING_SNAKE_CASE     )
														model.eval()

														def comm_check_on_output(lowerCamelCase_     ):
																					# let's still check that all the required stuff is there
																					self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None     )
																					self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None     )
																					self.parent.assertTrue(result.encoder_last_hidden_state is not None     )
																					# okay, now we need to check the logits shape
																					# due to the encoder compression, masks have a //4 spatial size
																					self.parent.assertEqual(
																					    result.masks_queries_logits.shape      ,					(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)      ,					)
																					# + 1 for null class
																					self.parent.assertEqual(
																					    result.class_queries_logits.shape      ,					(self.batch_size, self.num_queries, self.num_labels + 1)     )

														with torch.no_grad():
																					lowerCAmelCase__      =			model(pixel_values=_SCREAMING_SNAKE_CASE      ,					pixel_mask=_SCREAMING_SNAKE_CASE     )
																					lowerCAmelCase__      =			model(_SCREAMING_SNAKE_CASE     )

																					comm_check_on_output(_SCREAMING_SNAKE_CASE     )

																					lowerCAmelCase__      =			model(
																					    pixel_values=_SCREAMING_SNAKE_CASE      ,					pixel_mask=_SCREAMING_SNAKE_CASE      ,					mask_labels=_SCREAMING_SNAKE_CASE      ,					class_labels=_SCREAMING_SNAKE_CASE     )

														comm_check_on_output(_SCREAMING_SNAKE_CASE     )

														self.parent.assertTrue(result.loss is not None     )
														self.parent.assertEqual(result.loss.shape      ,					torch.Size([1]     )     )




@require_torch
class 	a__       (							a__       ,       a__       ,       unittest.TestCase						):





							'''simple docstring'''






							lowercase__    :     List[str]					=   (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
							lowercase__    :     Tuple					=   {"feature-extraction": MaskaFormerModel} if is_torch_available() else {}

							lowercase__    :     Any					=   False
							lowercase__    :     Tuple					=   False
							lowercase__    :     Optional[int]					=   False
							lowercase__    :     List[str]					=   False







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														lowerCAmelCase__      =			MaskaFormerModelTester(self     )
														lowerCAmelCase__      =			ConfigTester(self      ,					config_class=_SCREAMING_SNAKE_CASE      ,					has_text_modality=_SCREAMING_SNAKE_CASE     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														self.config_tester.run_common_tests()







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()
														self.model_tester.create_and_check_maskaformer_model(_SCREAMING_SNAKE_CASE      ,					**_SCREAMING_SNAKE_CASE      ,					output_hidden_states=_SCREAMING_SNAKE_CASE     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_SCREAMING_SNAKE_CASE     )







							@unittest.skip(reason='''Mask2Former does not use inputs_embeds'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
														pass







							@unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														pass







							@unittest.skip(reason='''Mask2Former is not a generative model'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
														pass







							@unittest.skip(reason='''Mask2Former does not use token embeddings'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														pass







							@require_torch_multi_gpu
							@unittest.skip(
							    reason='''Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							str:
														pass







							@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[int]:
														pass







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__      =			model_class(_SCREAMING_SNAKE_CASE     )
																					lowerCAmelCase__      =			inspect.signature(model.forward     )
																					# signature.parameters is an OrderedDict => so arg_names order is deterministic
																					lowerCAmelCase__      =			[*signature.parameters.keys()]

																					lowerCAmelCase__      =			['''pixel_values''']
																					self.assertListEqual(arg_names[:1]      ,					_SCREAMING_SNAKE_CASE     )







							@slow
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							str:






														for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
																					lowerCAmelCase__      =			MaskaFormerModel.from_pretrained(_SCREAMING_SNAKE_CASE     )
																					self.assertIsNotNone(_SCREAMING_SNAKE_CASE     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							int:
														lowerCAmelCase__      =			(self.model_tester.min_size,) * 2
														lowerCAmelCase__      =			{
														    '''pixel_values''': torch.randn((2, 3, *size)      ,					device=_SCREAMING_SNAKE_CASE     ),
														    '''mask_labels''': torch.randn((2, 10, *size)      ,					device=_SCREAMING_SNAKE_CASE     ),
														    '''class_labels''': torch.zeros(2      ,					10      ,					device=_SCREAMING_SNAKE_CASE     ).long(),
														}
														lowerCAmelCase__      =			self.model_tester.get_config()

														lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation(_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )
														lowerCAmelCase__      =			model(**_SCREAMING_SNAKE_CASE     )
														self.assertTrue(outputs.loss is not None     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
														lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()
														self.model_tester.create_and_check_maskaformer_model(_SCREAMING_SNAKE_CASE      ,					**_SCREAMING_SNAKE_CASE      ,					output_hidden_states=_SCREAMING_SNAKE_CASE     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Union[str, Any]:
														lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					lowerCAmelCase__      =			model_class(_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )
																					lowerCAmelCase__      =			model(**_SCREAMING_SNAKE_CASE      ,					output_attentions=_SCREAMING_SNAKE_CASE     )
																					self.assertTrue(outputs.attentions is not None     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:






														if not self.model_tester.is_training:
																					return

														lowerCAmelCase__      =			self.all_model_classes[1]
														lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()

														lowerCAmelCase__      =			model_class(_SCREAMING_SNAKE_CASE     )
														model.to(_SCREAMING_SNAKE_CASE     )
														model.train()

														lowerCAmelCase__      =			model(_SCREAMING_SNAKE_CASE      ,					mask_labels=_SCREAMING_SNAKE_CASE      ,					class_labels=_SCREAMING_SNAKE_CASE     ).loss
														loss.backward()







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							str:
														lowerCAmelCase__      =			self.all_model_classes[1]
														lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
														lowerCAmelCase__      =			True
														lowerCAmelCase__      =			True

														lowerCAmelCase__      =			model_class(_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )
														model.train()

														lowerCAmelCase__      =			model(_SCREAMING_SNAKE_CASE      ,					mask_labels=_SCREAMING_SNAKE_CASE      ,					class_labels=_SCREAMING_SNAKE_CASE     )

														lowerCAmelCase__      =			outputs.encoder_hidden_states[0]
														encoder_hidden_states.retain_grad()

														lowerCAmelCase__      =			outputs.pixel_decoder_hidden_states[0]
														pixel_decoder_hidden_states.retain_grad()

														lowerCAmelCase__      =			outputs.transformer_decoder_hidden_states[0]
														transformer_decoder_hidden_states.retain_grad()

														lowerCAmelCase__      =			outputs.attentions[0]
														attentions.retain_grad()

														outputs.loss.backward(retain_graph=_SCREAMING_SNAKE_CASE     )

														self.assertIsNotNone(encoder_hidden_states.grad     )
														self.assertIsNotNone(pixel_decoder_hidden_states.grad     )
														self.assertIsNotNone(transformer_decoder_hidden_states.grad     )
														self.assertIsNotNone(attentions.grad     )

__UpperCAmelCase					       =							1e-4


def       _snake_case     (  )		->					Dict:
							lowerCAmelCase__      =			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png''' )
							return image






@require_vision
@slow
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							@cached_property
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
														return "facebook/mask2former-swin-small-coco-instance"







							@cached_property
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
														return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints     ) if is_vision_available() else None







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[Any]:
														lowerCAmelCase__      =			MaskaFormerModel.from_pretrained(self.model_checkpoints     ).to(_SCREAMING_SNAKE_CASE     )
														lowerCAmelCase__      =			self.default_image_processor
														lowerCAmelCase__      =			prepare_img()
														lowerCAmelCase__      =			image_processor(_SCREAMING_SNAKE_CASE      ,					return_tensors='''pt'''     ).to(_SCREAMING_SNAKE_CASE     )
														lowerCAmelCase__      =			inputs['''pixel_values'''].shape
														# check size is divisible by 32
														self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
														# check size
														self.assertEqual(_SCREAMING_SNAKE_CASE      ,					(1, 3, 3_84, 3_84)     )

														with torch.no_grad():
																					lowerCAmelCase__      =			model(**_SCREAMING_SNAKE_CASE     )

														lowerCAmelCase__      =			torch.tensor(
														    [[-0.2_790, -1.0_717, -1.1_668], [-0.5_128, -0.3_128, -0.4_987], [-0.5_832, 0.1_971, -0.0_197]]     ).to(_SCREAMING_SNAKE_CASE     )
														self.assertTrue(
														    torch.allclose(
														        outputs.encoder_last_hidden_state[0, 0, :3, :3]      ,					_SCREAMING_SNAKE_CASE      ,					atol=_SCREAMING_SNAKE_CASE     )     )

														lowerCAmelCase__      =			torch.tensor(
														    [[0.8_973, 1.1_847, 1.1_776], [1.1_934, 1.5_040, 1.5_128], [1.1_153, 1.4_486, 1.4_951]]     ).to(_SCREAMING_SNAKE_CASE     )
														self.assertTrue(
														    torch.allclose(
														        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]      ,					_SCREAMING_SNAKE_CASE      ,					atol=_SCREAMING_SNAKE_CASE     )     )

														lowerCAmelCase__      =			torch.tensor(
														    [[2.1_152, 1.7_000, -0.8_603], [1.5_808, 1.8_004, -0.9_353], [1.6_043, 1.7_495, -0.5_999]]     ).to(_SCREAMING_SNAKE_CASE     )
														self.assertTrue(
														    torch.allclose(
														        outputs.transformer_decoder_last_hidden_state[0, :3, :3]      ,					_SCREAMING_SNAKE_CASE      ,					atol=_SCREAMING_SNAKE_CASE     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Any:
														lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(_SCREAMING_SNAKE_CASE     ).eval()
														lowerCAmelCase__      =			self.default_image_processor
														lowerCAmelCase__      =			prepare_img()
														lowerCAmelCase__      =			image_processor(_SCREAMING_SNAKE_CASE      ,					return_tensors='''pt'''     ).to(_SCREAMING_SNAKE_CASE     )
														lowerCAmelCase__      =			inputs['''pixel_values'''].shape
														# check size is divisible by 32
														self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0     )
														# check size
														self.assertEqual(_SCREAMING_SNAKE_CASE      ,					(1, 3, 3_84, 3_84)     )

														with torch.no_grad():
																					lowerCAmelCase__      =			model(**_SCREAMING_SNAKE_CASE     )
														# masks_queries_logits
														lowerCAmelCase__      =			outputs.masks_queries_logits
														self.assertEqual(
														    masks_queries_logits.shape      ,					(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)     )
														lowerCAmelCase__      =			[
														    [-8.7_839, -9.0_056, -8.8_121],
														    [-7.4_104, -7.0_313, -6.5_401],
														    [-6.6_105, -6.3_427, -6.4_675],
														]
														lowerCAmelCase__      =			torch.tensor(_SCREAMING_SNAKE_CASE     ).to(_SCREAMING_SNAKE_CASE     )
														self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]      ,					_SCREAMING_SNAKE_CASE      ,					atol=_SCREAMING_SNAKE_CASE     )     )
														# class_queries_logits
														lowerCAmelCase__      =			outputs.class_queries_logits
														self.assertEqual(class_queries_logits.shape      ,					(1, model.config.num_queries, model.config.num_labels + 1)     )
														lowerCAmelCase__      =			torch.tensor(
														    [
														        [1.8_324, -8.0_835, -4.1_922],
														        [0.8_450, -9.0_050, -3.6_053],
														        [0.3_045, -7.7_293, -3.0_275],
														    ]     ).to(_SCREAMING_SNAKE_CASE     )
														self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]      ,					_SCREAMING_SNAKE_CASE      ,					atol=_SCREAMING_SNAKE_CASE     )     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														lowerCAmelCase__      =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(_SCREAMING_SNAKE_CASE     ).eval()
														lowerCAmelCase__      =			self.default_image_processor

														lowerCAmelCase__      =			image_processor(
														    [np.zeros((3, 8_00, 13_33)     ), np.zeros((3, 8_00, 13_33)     )]      ,					segmentation_maps=[np.zeros((3_84, 3_84)     ).astype(np.floataa     ), np.zeros((3_84, 3_84)     ).astype(np.floataa     )]      ,					return_tensors='''pt'''      ,					)

														lowerCAmelCase__      =			inputs['''pixel_values'''].to(_SCREAMING_SNAKE_CASE     )
														lowerCAmelCase__      =			[el.to(_SCREAMING_SNAKE_CASE     ) for el in inputs['''mask_labels''']]
														lowerCAmelCase__      =			[el.to(_SCREAMING_SNAKE_CASE     ) for el in inputs['''class_labels''']]

														with torch.no_grad():
																					lowerCAmelCase__      =			model(**_SCREAMING_SNAKE_CASE     )

														self.assertTrue(outputs.loss is not None     )



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


__UpperCAmelCase					       =							logging.getLogger(__name__)




class 	a__       (							a__						):





							'''simple docstring'''






							def __init__(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_=None     ) ->							Dict:
														super().__init__(
														    lowerCamelCase_      ,					question_encoder_tokenizer=lowerCamelCase_      ,					generator_tokenizer=lowerCamelCase_      ,					index=lowerCamelCase_      ,					init_retrieval=lowerCamelCase_      ,					)
														lowerCAmelCase__      =			None







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_     ) ->							str:
														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=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      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_=torch.floataa     ) ->							Union[str, Any]:
														lowerCAmelCase__      =			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     ) ->							str:
														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'''     ))      ,					lowerCamelCase_     )
														return ifname







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

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

														# 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(lowerCamelCase_     )]
														dist.gather(torch.tensor(lowerCamelCase_     )      ,					dst=0      ,					gather_list=lowerCamelCase_      ,					group=self.process_group     )

														# scatter logic
														lowerCAmelCase__      =			question_hidden_states.shape[0]
														lowerCAmelCase__      =			[]
														lowerCAmelCase__      =			[]
														if self._is_main():
																					assert len(lowerCamelCase_     ) == world_size
																					lowerCAmelCase__     ,      lowerCAmelCase__      =			self._main_retrieve(torch.cat(lowerCamelCase_     ).numpy()      ,					lowerCamelCase_     )
																					lowerCAmelCase__     ,      lowerCAmelCase__      =			torch.tensor(lowerCamelCase_     ), torch.tensor(lowerCamelCase_     )
																					lowerCAmelCase__      =			self._chunk_tensor(lowerCamelCase_      ,					lowerCamelCase_     )
																					lowerCAmelCase__      =			self._chunk_tensor(lowerCamelCase_      ,					lowerCamelCase_     )
														lowerCAmelCase__      =			self._scattered(lowerCamelCase_      ,					[n_queries, n_docs]      ,					target_type=torch.intaa     )
														lowerCAmelCase__      =			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_     )