



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				:

			def __init__(self,   SCREAMING_SNAKE_CASE_,   )       ->   int:
						UpperCAmelCase_:		Any							      =     parent
						UpperCAmelCase_:		int							      =     13
						UpperCAmelCase_:		Union[str, Any]							      =     7
						UpperCAmelCase_:		Optional[Any]							      =     True
						UpperCAmelCase_:		int							      =     True
						UpperCAmelCase_:		Optional[Any]							      =     True
						UpperCAmelCase_:		Union[str, Any]							      =     99
						UpperCAmelCase_:		Optional[int]							      =     32
						UpperCAmelCase_:		Tuple							      =     2
						UpperCAmelCase_:		List[str]							      =     4
						UpperCAmelCase_:		Tuple							      =     37
						UpperCAmelCase_:		List[Any]							      =     """gelu"""
						UpperCAmelCase_:		Tuple							      =     0.1
						UpperCAmelCase_:		Tuple							      =     0.1
						UpperCAmelCase_:		Tuple							      =     512
						UpperCAmelCase_:		Dict							      =     16
						UpperCAmelCase_:		int							      =     2
						UpperCAmelCase_:		Any							      =     0.0_2
						UpperCAmelCase_:		int							      =     3
						UpperCAmelCase_:		int							      =     4
						UpperCAmelCase_:		Any							      =     None

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		List[str]							      =     ids_tensor([self.batch_size, self.seq_length],   self.vocab_size					)

						UpperCAmelCase_:		Any							      =     None
						if self.use_input_mask:
									UpperCAmelCase_:		List[str]							      =     random_attention_mask([self.batch_size, self.seq_length]					)

						UpperCAmelCase_:		List[Any]							      =     None
						UpperCAmelCase_:		Dict							      =     None
						UpperCAmelCase_:		Optional[Any]							      =     None
						if self.use_labels:
									UpperCAmelCase_:		Dict							      =     ids_tensor([self.batch_size],   self.type_sequence_label_size					)
									UpperCAmelCase_:		int							      =     ids_tensor([self.batch_size, self.seq_length],   self.num_labels					)
									UpperCAmelCase_:		List[str]							      =     ids_tensor([self.batch_size],   self.num_choices					)

						UpperCAmelCase_:		str							      =     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      __snake_case       (self					)       ->   int:
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		int							      =     self.prepare_config_and_inputs()

						UpperCAmelCase_:		Any							      =     True
						UpperCAmelCase_:		Optional[Any]							      =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
						UpperCAmelCase_:		int							      =     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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						UpperCAmelCase_:		List[Any]							      =     TFEsmModel(config=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		str							      =     {"""input_ids""": input_ids, """attention_mask""": input_mask}
						UpperCAmelCase_:		int							      =     model(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		str							      =     [input_ids, input_mask]
						UpperCAmelCase_:		Tuple							      =     model(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		int							      =     model(SCREAMING_SNAKE_CASE_					)

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

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   )       ->   Optional[Any]:
						UpperCAmelCase_:		str							      =     True

						UpperCAmelCase_:		int							      =     TFEsmModel(config=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		str							      =     {
						    """input_ids""": input_ids,
						    """attention_mask""": input_mask,
						    """encoder_hidden_states""": encoder_hidden_states,
						    """encoder_attention_mask""": encoder_attention_mask,
						}
						UpperCAmelCase_:		Union[str, Any]							      =     model(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Union[str, Any]							      =     [input_ids, input_mask]
						UpperCAmelCase_:		Any							      =     model(SCREAMING_SNAKE_CASE_,   encoder_hidden_states=SCREAMING_SNAKE_CASE_					)

						# Also check the case where encoder outputs are not passed
						UpperCAmelCase_:		Any							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_					)

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

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						UpperCAmelCase_:		Dict							      =     TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		int							      =     model([input_ids, input_mask]					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.vocab_size)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Union[str, Any]							      =     self.num_labels
						UpperCAmelCase_:		List[str]							      =     TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[str]							      =     {"""input_ids""": input_ids, """attention_mask""": input_mask}
						UpperCAmelCase_:		Dict							      =     model(SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.num_labels)					)

			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		List[str]							      =     self.prepare_config_and_inputs()
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		List[str]							      =     config_and_inputs
						UpperCAmelCase_:		List[Any]							      =     {"""input_ids""": input_ids, """attention_mask""": input_mask}
						return config, inputs_dict







@require_tf
class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			A      		=	(
			    {
			        '''feature-extraction''': TFEsmModel,
			        '''fill-mask''': TFEsmForMaskedLM,
			        '''text-classification''': TFEsmForSequenceClassification,
			        '''token-classification''': TFEsmForTokenClassification,
			        '''zero-shot''': TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			A      		=	False
			A      		=	False

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Union[str, Any]							      =     TFEsmModelTester(self					)
						UpperCAmelCase_:		int							      =     ConfigTester(self,   config_class=SCREAMING_SNAKE_CASE_,   hidden_size=37					)

			def      __snake_case       (self					)       ->   int:
						self.config_tester.run_common_tests()

			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		Any							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Tuple:
						UpperCAmelCase_:		Optional[int]							      =     self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		int							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   List[str]:
						UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_					)

			@slow
			def      __snake_case       (self					)       ->   Any:
						for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase_:		List[str]							      =     TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_					)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_					)

			@unittest.skip("""Protein models do not support embedding resizing."""					)
			def      __snake_case       (self					)       ->   Dict:
						pass

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

			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_      ,							UpperCAmelCase_:		Dict							      =     self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									UpperCAmelCase_:		List[Any]							      =     model_class(SCREAMING_SNAKE_CASE_					)
									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
												UpperCAmelCase_:		Dict							      =     model.get_bias()
												assert isinstance(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)
												for k, v in name.items():
															assert isinstance(SCREAMING_SNAKE_CASE_,   tf.Variable					)
									else:
												UpperCAmelCase_:		Union[str, Any]							      =     model.get_output_embeddings()
												assert x is None
												UpperCAmelCase_:		str							      =     model.get_bias()
												assert name is None







@require_tf
class 				_a				(   unittest.TestCase					):

			@slow
			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		List[str]							      =     TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

						UpperCAmelCase_:		int							      =     tf.constant([[0, 1, 2, 3, 4, 5]]					)
						UpperCAmelCase_:		Optional[Any]							      =     model(SCREAMING_SNAKE_CASE_					)[0]
						UpperCAmelCase_:		str							      =     [1, 6, 33]
						self.assertEqual(list(output.numpy().shape					),   SCREAMING_SNAKE_CASE_					)
						# compare the actual values for a slice.
						UpperCAmelCase_:		List[str]							      =     tf.constant(
						    [
						        [
						            [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
						            [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
						            [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
						        ]
						    ]					)
						self.assertTrue(numpy.allclose(output[:, :3, :3].numpy(),   expected_slice.numpy(),   atol=1E-2					)					)

			@slow
			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		Optional[int]							      =     TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

						UpperCAmelCase_:		List[Any]							      =     tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
						UpperCAmelCase_:		Tuple							      =     model(SCREAMING_SNAKE_CASE_					)[0]
						# compare the actual values for a slice.
						UpperCAmelCase_:		str							      =     tf.constant(
						    [
						        [
						            [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
						            [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
						            [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
						        ]
						    ]					)
						self.assertTrue(numpy.allclose(output[:, :3, :3].numpy(),   expected_slice.numpy(),   atol=1E-4					)					)











#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist







def  lowerCAmelCase_	(*lowerCAmelCase__:		Dict  ):





			"""simple docstring"""


			with open(lowerCAmelCase__		,						"""r"""  ) as fh:
						fcntl.flock(lowerCAmelCase__		,						fcntl.LOCK_EX  )
						try:
									print(*lowerCAmelCase__  )
						finally:
									fcntl.flock(lowerCAmelCase__		,						fcntl.LOCK_UN  )


a      :  Tuple			     =			int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
a      :  Union[str, Any]			     =			torch.device('cuda', local_rank)
a      :  Any			     =			socket.gethostname()

a      :  Optional[Any]			     =			F'''[{hostname}-{local_rank}]'''

try:
	# test distributed
	dist.init_process_group('nccl')
	dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
	dist.barrier()

	# test cuda is available and can allocate memory
	torch.cuda.is_available()
	torch.ones(1).cuda(local_rank)

	# global rank
	a      :  Optional[Any]			     =			dist.get_rank()
	a      :  List[Any]			     =			dist.get_world_size()

	printflock(F'''{gpu} is OK (global rank: {rank}/{world_size})''')

	dist.barrier()
	if rank == 0:
		printflock(F'''pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}''')

except Exception:
	printflock(F'''{gpu} is broken''')
	raise






