


'''simple docstring'''






from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers.models.esm.modeling_tf_esm import (
		    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFEsmForMaskedLM,
		    TFEsmForSequenceClassification,
		    TFEsmForTokenClassification,
		    TFEsmModel,
		)




class 	a__       :





							'''simple docstring'''






							def __init__(	self      ,					lowerCamelCase_      ,					) ->							Any:
														lowerCAmelCase__      =			parent
														lowerCAmelCase__      =			13
														lowerCAmelCase__      =			7
														lowerCAmelCase__      =			True
														lowerCAmelCase__      =			True
														lowerCAmelCase__      =			True
														lowerCAmelCase__      =			99
														lowerCAmelCase__      =			32
														lowerCAmelCase__      =			2
														lowerCAmelCase__      =			4
														lowerCAmelCase__      =			37
														lowerCAmelCase__      =			'''gelu'''
														lowerCAmelCase__      =			0.1
														lowerCAmelCase__      =			0.1
														lowerCAmelCase__      =			5_12
														lowerCAmelCase__      =			16
														lowerCAmelCase__      =			2
														lowerCAmelCase__      =			0.02
														lowerCAmelCase__      =			3
														lowerCAmelCase__      =			4
														lowerCAmelCase__      =			None







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size     )

														lowerCAmelCase__      =			None
														if self.use_input_mask:
																					lowerCAmelCase__      =			random_attention_mask([self.batch_size, self.seq_length]     )

														lowerCAmelCase__      =			None
														lowerCAmelCase__      =			None
														lowerCAmelCase__      =			None
														if self.use_labels:
																					lowerCAmelCase__      =			ids_tensor([self.batch_size]      ,					self.type_sequence_label_size     )
																					lowerCAmelCase__      =			ids_tensor([self.batch_size, self.seq_length]      ,					self.num_labels     )
																					lowerCAmelCase__      =			ids_tensor([self.batch_size]      ,					self.num_choices     )

														lowerCAmelCase__      =			EsmConfig(
														    vocab_size=self.vocab_size      ,					hidden_size=self.hidden_size      ,					num_hidden_layers=self.num_hidden_layers      ,					pad_token_id=1      ,					num_attention_heads=self.num_attention_heads      ,					intermediate_size=self.intermediate_size      ,					hidden_act=self.hidden_act      ,					hidden_dropout_prob=self.hidden_dropout_prob      ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,					max_position_embeddings=self.max_position_embeddings      ,					type_vocab_size=self.type_vocab_size      ,					initializer_range=self.initializer_range      ,					)

														return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Optional[int]:
														(
														    (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      
														)      =			self.prepare_config_and_inputs()

														lowerCAmelCase__      =			True
														lowerCAmelCase__      =			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
														lowerCAmelCase__      =			ids_tensor([self.batch_size, self.seq_length]      ,					vocab_size=2     )

														return (
														    config,
														    input_ids,
														    input_mask,
														    sequence_labels,
														    token_labels,
														    choice_labels,
														    encoder_hidden_states,
														    encoder_attention_mask,
														)







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							str:
														lowerCAmelCase__      =			TFEsmModel(config=lowerCamelCase_     )
														lowerCAmelCase__      =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														lowerCAmelCase__      =			model(lowerCamelCase_     )

														lowerCAmelCase__      =			[input_ids, input_mask]
														lowerCAmelCase__      =			model(lowerCamelCase_     )

														lowerCAmelCase__      =			model(lowerCamelCase_     )

														self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size)     )







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

														lowerCAmelCase__      =			TFEsmModel(config=lowerCamelCase_     )
														lowerCAmelCase__      =			{
														    '''input_ids''': input_ids,
														    '''attention_mask''': input_mask,
														    '''encoder_hidden_states''': encoder_hidden_states,
														    '''encoder_attention_mask''': encoder_attention_mask,
														}
														lowerCAmelCase__      =			model(lowerCamelCase_     )

														lowerCAmelCase__      =			[input_ids, input_mask]
														lowerCAmelCase__      =			model(lowerCamelCase_      ,					encoder_hidden_states=lowerCamelCase_     )

														# Also check the case where encoder outputs are not passed
														lowerCAmelCase__      =			model(lowerCamelCase_      ,					attention_mask=lowerCamelCase_     )

														self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size)     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							Dict:
														lowerCAmelCase__      =			TFEsmForMaskedLM(config=lowerCamelCase_     )
														lowerCAmelCase__      =			model([input_ids, input_mask]     )
														self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.vocab_size)     )







							def 						__SCREAMING_SNAKE_CASE							(	self      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_      ,					lowerCamelCase_     ) ->							str:
														lowerCAmelCase__      =			self.num_labels
														lowerCAmelCase__      =			TFEsmForTokenClassification(config=lowerCamelCase_     )
														lowerCAmelCase__      =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														lowerCAmelCase__      =			model(lowerCamelCase_     )
														self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.num_labels)     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			self.prepare_config_and_inputs()
														(
														    (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      (
														    lowerCAmelCase__
														)     ,      
														)      =			config_and_inputs
														lowerCAmelCase__      =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														return config, inputs_dict




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





							'''simple docstring'''






							lowercase__    :     Any					=   (
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							lowercase__    :     Dict					=   (
							    {
							        "feature-extraction": TFEsmModel,
							        "fill-mask": TFEsmForMaskedLM,
							        "text-classification": TFEsmForSequenceClassification,
							        "token-classification": TFEsmForTokenClassification,
							        "zero-shot": TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							lowercase__    :     List[str]					=   False
							lowercase__    :     Union[str, Any]					=   False







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
														lowerCAmelCase__      =			TFEsmModelTester(self     )
														lowerCAmelCase__      =			ConfigTester(self      ,					config_class=lowerCamelCase_      ,					hidden_size=37     )







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







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[Any]:
														lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs_for_decoder()
														self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase_     )







							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
														lowerCAmelCase__      =			self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_token_classification(*lowerCamelCase_     )







							@slow
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					lowerCAmelCase__      =			TFEsmModel.from_pretrained(lowerCamelCase_     )
																					self.assertIsNotNone(lowerCamelCase_     )







							@unittest.skip('''Protein models do not support embedding resizing.'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							int:
														pass







							@unittest.skip('''Protein models do not support embedding resizing.'''     )
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														pass







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

														for model_class in self.all_model_classes:
																					lowerCAmelCase__      =			model_class(lowerCamelCase_     )
																					assert isinstance(model.get_input_embeddings()      ,					tf.keras.layers.Layer     )
																					if model_class is TFEsmForMaskedLM:
																												# Output embedding test differs from the main test because they're a matrix, not a layer
																												lowerCAmelCase__      =			model.get_bias()
																												assert isinstance(lowerCamelCase_      ,					lowerCamelCase_     )
																												for k, v in name.items():
																																			assert isinstance(lowerCamelCase_      ,					tf.Variable     )
																					else:
																												lowerCAmelCase__      =			model.get_output_embeddings()
																												assert x is None
																												lowerCAmelCase__      =			model.get_bias()
																												assert name is None




@require_tf
class 	a__       (							unittest.TestCase						):





							'''simple docstring'''






							@slow
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Dict:
														lowerCAmelCase__      =			TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''     )

														lowerCAmelCase__      =			tf.constant([[0, 1, 2, 3, 4, 5]]     )
														lowerCAmelCase__      =			model(lowerCamelCase_     )[0]
														lowerCAmelCase__      =			[1, 6, 33]
														self.assertEqual(list(output.numpy().shape     )      ,					lowerCamelCase_     )
														# compare the actual values for a slice.
														lowerCAmelCase__      =			tf.constant(
														    [
														        [
														            [8.921_518, -10.589_814, -6.4_671_307],
														            [-6.3_967_156, -13.911_377, -1.1_211_915],
														            [-7.781_247, -13.951_557, -3.740_592],
														        ]
														    ]     )
														self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,					expected_slice.numpy()      ,					atol=1e-2     )     )







							@slow
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							Tuple:
														lowerCAmelCase__      =			TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''     )

														lowerCAmelCase__      =			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]     )
														lowerCAmelCase__      =			model(lowerCamelCase_     )[0]
														# compare the actual values for a slice.
														lowerCAmelCase__      =			tf.constant(
														    [
														        [
														            [0.14_443_092, 0.54_125_327, 0.3_247_739],
														            [0.30_340_484, 0.00_526_676, 0.31_077_722],
														            [0.32_278_043, -0.24_987_096, 0.3_414_628],
														        ]
														    ]     )
														self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,					expected_slice.numpy()      ,					atol=1e-4     )     )



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