"""simple docstring"""

from abc import ABC, abstractmethod
from typing import List, Optional







class      UpperCAmelCase_       (						_lowercase):






							def __init__(		self :      Tuple  )       ->       Tuple:
												self.test()






							def       _UpperCamelCase			(		self :      str  )       ->       int:
												_UpperCamelCase      						=	0
												_UpperCamelCase      						=	False
												while not completed:
																	if counter == 1:
																						self.reset()
																	_UpperCamelCase      						=	self.advance()
																	if not self.does_advance(__UpperCamelCase  ):
																						raise Exception(
																						    '''Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.'''  )

																	_UpperCamelCase						,							_UpperCamelCase						,							_UpperCamelCase      						=	self.update(__UpperCamelCase  )
																	counter += 1

																	if counter > 1_0000:
																						raise Exception('''update() does not fulfill the constraint.'''  )

												if self.remaining() != 0:
																	raise Exception('''Custom Constraint is not defined correctly.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      str  )       ->       Optional[int]:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      List[str]     ,						__UpperCamelCase :      Any  )       ->       Union[str, Any]:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      Dict     ,						__UpperCamelCase :      Dict  )       ->       List[Any]:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      int  )       ->       Tuple:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      Dict  )       ->       Any:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )






							@abstractmethod
							def       _UpperCamelCase			(		self :      Dict     ,						__UpperCamelCase :      Optional[Any]=False  )       ->       Any:
												raise NotImplementedError(
												    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.'''  )







class      UpperCAmelCase_       (						_lowercase):






							def __init__(		self :      List[Any]     ,						__UpperCamelCase :      List[str]  )       ->       List[str]:
												super(__UpperCamelCase     ,						self  ).__init__()

												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ) or len(__UpperCamelCase  ) == 0:
																	raise ValueError(F'''`token_ids` has to be a non-empty list, but is {token_ids}.'''  )
												if any((not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ) or token_id < 0) for token_id in token_ids  ):
																	raise ValueError(F'''Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.'''  )

												_UpperCamelCase      						=	token_ids

												_UpperCamelCase      						=	len(self.token_ids  )
												_UpperCamelCase      						=	-1  # the index of the currently fulfilled step
												_UpperCamelCase      						=	False






							def       _UpperCamelCase			(		self :      str  )       ->       List[str]:
												if self.completed:
																	return None
												return self.token_ids[self.fulfilled_idx + 1]






							def       _UpperCamelCase			(		self :      int     ,						__UpperCamelCase :      List[Any]  )       ->       Optional[int]:
												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(__UpperCamelCase  )}'''  )

												if self.completed:
																	return False

												return token_id == self.token_ids[self.fulfilled_idx + 1]






							def       _UpperCamelCase			(		self :      List[Any]     ,						__UpperCamelCase :      Optional[Any]  )       ->       Tuple:
												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(__UpperCamelCase  )}'''  )

												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	False

												if self.does_advance(__UpperCamelCase  ):
																	self.fulfilled_idx += 1
																	_UpperCamelCase      						=	True
																	if self.fulfilled_idx == (self.seqlen - 1):
																						_UpperCamelCase      						=	True
																	_UpperCamelCase      						=	completed
												else:
																	# failed to make progress.
																	_UpperCamelCase      						=	True
																	self.reset()
												return stepped, completed, reset






							def       _UpperCamelCase			(		self :      int  )       ->       Union[str, Any]:
												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	0






							def       _UpperCamelCase			(		self :      Optional[Any]  )       ->       Optional[Any]:
												return self.seqlen - (self.fulfilled_idx + 1)






							def       _UpperCamelCase			(		self :      Any     ,						__UpperCamelCase :      int=False  )       ->       Dict:
												_UpperCamelCase      						=	PhrasalConstraint(self.token_ids  )

												if stateful:
																	_UpperCamelCase      						=	self.seqlen
																	_UpperCamelCase      						=	self.fulfilled_idx
																	_UpperCamelCase      						=	self.completed

												return new_constraint







class      UpperCAmelCase_       :






							def __init__(		self :      Dict     ,						__UpperCamelCase :      List[str]     ,						__UpperCamelCase :      Dict=True  )       ->       Tuple:
												_UpperCamelCase      						=	max([len(__UpperCamelCase  ) for one in nested_token_ids]  )

												_UpperCamelCase      						=	{}
												for token_ids in nested_token_ids:
																	_UpperCamelCase      						=	root
																	for tidx, token_id in enumerate(__UpperCamelCase  ):
																						if token_id not in level:
																											_UpperCamelCase      						=	{}

																						_UpperCamelCase      						=	level[token_id]

												if no_subsets and self.has_subsets(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(
																	    '''Each list in `nested_token_ids` can\'t be a complete subset of another list, but is'''
																	    F''' {nested_token_ids}.'''  )

												_UpperCamelCase      						=	root






							def       _UpperCamelCase			(		self :      Optional[Any]     ,						__UpperCamelCase :      Dict  )       ->       Any:
												_UpperCamelCase      						=	self.trie

												for current_token in current_seq:
																	_UpperCamelCase      						=	start[current_token]

												_UpperCamelCase      						=	list(start.keys()  )

												return next_tokens






							def       _UpperCamelCase			(		self :      int     ,						__UpperCamelCase :      int  )       ->       Optional[Any]:
												_UpperCamelCase      						=	self.next_tokens(__UpperCamelCase  )

												return len(__UpperCamelCase  ) == 0






							def       _UpperCamelCase			(		self :      Any     ,						__UpperCamelCase :      Tuple  )       ->       List[str]:
												_UpperCamelCase      						=	list(root.values()  )
												if len(__UpperCamelCase  ) == 0:
																	return 1
												else:
																	return sum([self.count_leaves(__UpperCamelCase  ) for nn in next_nodes]  )






							def       _UpperCamelCase			(		self :      Union[str, Any]     ,						__UpperCamelCase :      List[Any]     ,						__UpperCamelCase :      Optional[Any]  )       ->       int:
												_UpperCamelCase      						=	self.count_leaves(__UpperCamelCase  )
												return len(__UpperCamelCase  ) != leaf_count







class      UpperCAmelCase_       (						_lowercase):






							def __init__(		self :      Optional[int]     ,						__UpperCamelCase :      Dict  )       ->       str:
												super(__UpperCamelCase     ,						self  ).__init__()

												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ) or len(__UpperCamelCase  ) == 0:
																	raise ValueError(F'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.'''  )
												if any(not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ) for token_ids in nested_token_ids  ):
																	raise ValueError(F'''`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.'''  )
												if any(
												    any((not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ) or token_id < 0) for token_id in token_ids  )
												    for token_ids in nested_token_ids  ):
																	raise ValueError(
																	    F'''Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}.'''  )

												_UpperCamelCase      						=	DisjunctiveTrie(__UpperCamelCase  )
												_UpperCamelCase      						=	nested_token_ids

												_UpperCamelCase      						=	self.trie.max_height
												_UpperCamelCase      						=	[]
												_UpperCamelCase      						=	False






							def       _UpperCamelCase			(		self :      Any  )       ->       Tuple:
												_UpperCamelCase      						=	self.trie.next_tokens(self.current_seq  )

												if len(__UpperCamelCase  ) == 0:
																	return None
												else:
																	return token_list






							def       _UpperCamelCase			(		self :      List[str]     ,						__UpperCamelCase :      Union[str, Any]  )       ->       Tuple:
												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__UpperCamelCase  )}'''  )

												_UpperCamelCase      						=	self.trie.next_tokens(self.current_seq  )

												return token_id in next_tokens






							def       _UpperCamelCase			(		self :      Optional[int]     ,						__UpperCamelCase :      Any  )       ->       List[str]:
												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(__UpperCamelCase  )}'''  )

												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	False

												if self.does_advance(__UpperCamelCase  ):
																	self.current_seq.append(__UpperCamelCase  )
																	_UpperCamelCase      						=	True
												else:
																	_UpperCamelCase      						=	True
																	self.reset()

												_UpperCamelCase      						=	self.trie.reached_leaf(self.current_seq  )
												_UpperCamelCase      						=	completed

												return stepped, completed, reset






							def       _UpperCamelCase			(		self :      Union[str, Any]  )       ->       Union[str, Any]:
												_UpperCamelCase      						=	False
												_UpperCamelCase      						=	[]






							def       _UpperCamelCase			(		self :      List[str]  )       ->       List[Any]:
												if self.completed:
																	# since this can be completed without reaching max height
																	return 0
												else:
																	return self.seqlen - len(self.current_seq  )






							def       _UpperCamelCase			(		self :      str     ,						__UpperCamelCase :      int=False  )       ->       List[Any]:
												_UpperCamelCase      						=	DisjunctiveConstraint(self.token_ids  )

												if stateful:
																	_UpperCamelCase      						=	self.seqlen
																	_UpperCamelCase      						=	self.current_seq
																	_UpperCamelCase      						=	self.completed

												return new_constraint







class      UpperCAmelCase_       :






							def __init__(		self :      Any     ,						__UpperCamelCase :      str  )       ->       Dict:
												_UpperCamelCase      						=	constraints

												# max # of steps required to fulfill a given constraint
												_UpperCamelCase      						=	max([c.seqlen for c in constraints]  )
												_UpperCamelCase      						=	len(__UpperCamelCase  )
												_UpperCamelCase      						=	False

												self.init_state()






							def       _UpperCamelCase			(		self :      Optional[int]  )       ->       Union[str, Any]:
												_UpperCamelCase      						=	[]
												_UpperCamelCase      						=	None
												_UpperCamelCase      						=	[constraint.copy(stateful=__UpperCamelCase  ) for constraint in self.constraints]






							def       _UpperCamelCase			(		self :      int  )       ->       List[str]:
												_UpperCamelCase      						=	0
												if self.inprogress_constraint:
																	# extra points for having a constraint mid-fulfilled
																	add += self.max_seqlen - self.inprogress_constraint.remaining()

												return (len(self.complete_constraints  ) * self.max_seqlen) + add






							def       _UpperCamelCase			(		self :      str  )       ->       List[str]:
												_UpperCamelCase      						=	[]
												if self.inprogress_constraint is None:
																	for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
																						_UpperCamelCase      						=	constraint.advance()
																						if isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																											token_list.append(__UpperCamelCase  )
																						elif isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																											token_list.extend(__UpperCamelCase  )
												else:
																	_UpperCamelCase      						=	self.inprogress_constraint.advance()
																	if isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																						token_list.append(__UpperCamelCase  )
																	elif isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																						token_list.extend(__UpperCamelCase  )

												if len(__UpperCamelCase  ) == 0:
																	return None
												else:
																	return token_list






							def       _UpperCamelCase			(		self :      Optional[int]     ,						__UpperCamelCase :      Dict  )       ->       Dict:
												self.init_state()

												if token_ids is not None:
																	for token in token_ids:
																						# completes or steps **one** constraint
																						_UpperCamelCase						,							_UpperCamelCase      						=	self.add(__UpperCamelCase  )

																						# the entire list of constraints are fulfilled
																						if self.completed:
																											break






							def       _UpperCamelCase			(		self :      Optional[int]     ,						__UpperCamelCase :      Union[str, Any]  )       ->       List[str]:
												if not isinstance(__UpperCamelCase     ,						__UpperCamelCase  ):
																	raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.'''  )

												_UpperCamelCase						,							_UpperCamelCase      						=	False, False

												if self.completed:
																	_UpperCamelCase      						=	True
																	_UpperCamelCase      						=	False
																	return complete, stepped

												if self.inprogress_constraint is not None:
																	# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
																	# job, simply update the state

																	_UpperCamelCase						,							_UpperCamelCase						,							_UpperCamelCase      						=	self.inprogress_constraint.update(__UpperCamelCase  )
																	if reset:
																						# 1. If the next token breaks the progress, then we must restart.
																						#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

																						#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
																						#     constraint, not the full list of constraints.

																						self.pending_constraints.append(self.inprogress_constraint.copy(stateful=__UpperCamelCase  )  )
																						_UpperCamelCase      						=	None

																	if complete:
																						# 2. If the next token completes the constraint, move it to completed list, set
																						#     inprogress to None. If there are no pending constraints either, then this full list of constraints
																						#     is complete.

																						self.complete_constraints.append(self.inprogress_constraint  )
																						_UpperCamelCase      						=	None

																						if len(self.pending_constraints  ) == 0:
																											# we're done!
																											_UpperCamelCase      						=	True

												else:
																	# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
																	# of constraints?

																	for cidx, pending_constraint in enumerate(self.pending_constraints  ):
																						if pending_constraint.does_advance(__UpperCamelCase  ):
																											_UpperCamelCase						,							_UpperCamelCase						,							_UpperCamelCase      						=	pending_constraint.update(__UpperCamelCase  )

																											if not stepped:
																																raise Exception(
																																    '''`constraint.update(token_id)` is not yielding incremental progress, '''
																																    '''even though `constraint.does_advance(token_id)` is true.'''  )

																											if complete:
																																self.complete_constraints.append(__UpperCamelCase  )
																																_UpperCamelCase      						=	None

																											if not complete and stepped:
																																_UpperCamelCase      						=	pending_constraint

																											if complete or stepped:
																																# If we made any progress at all, then it's at least not a "pending constraint".

																																_UpperCamelCase      						=	(
																																    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
																																)

																																if len(self.pending_constraints  ) == 0 and self.inprogress_constraint is None:
																																					# If there's no longer any pending after this and no inprogress either, then we must be
																																					# complete.

																																					_UpperCamelCase      						=	True

																																break  # prevent accidentally stepping through multiple constraints with just one token.

												return complete, stepped






							def       _UpperCamelCase			(		self :      Optional[int]     ,						__UpperCamelCase :      str=True  )       ->       Union[str, Any]:
												_UpperCamelCase      						=	ConstraintListState(self.constraints  )  # we actually never though self.constraints objects
												# throughout this process. So it's at initialization state.

												if stateful:
																	_UpperCamelCase      						=	[
																	    constraint.copy(stateful=__UpperCamelCase  ) for constraint in self.complete_constraints
																	]
																	if self.inprogress_constraint is not None:
																						_UpperCamelCase      						=	self.inprogress_constraint.copy(stateful=__UpperCamelCase  )
																	_UpperCamelCase      						=	[constraint.copy() for constraint in self.pending_constraints]

												return new_state




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





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase			,  ):







   """simple docstring"""
   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_      =  512
   lowerCamelCase_      =  16
   lowerCamelCase_      =  2
   lowerCamelCase_      =  0.02
   lowerCamelCase_      =  3
   lowerCamelCase_      =  4
   lowerCamelCase_      =  None






  def   snake_case		(    self     ):







   """simple docstring"""
   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   snake_case		(    self     ):







   """simple docstring"""
   (
       (
       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   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  TFEsmModel(config=UpperCamelCase     )
   lowerCamelCase_      =  {"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase_      =  model(UpperCamelCase     )

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

   lowerCamelCase_      =  model(UpperCamelCase     )

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






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  ):







   """simple docstring"""
   lowerCamelCase_      =  True

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

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

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

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






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  TFEsmForMaskedLM(config=UpperCamelCase     )
   lowerCamelCase_      =  model([input_ids, input_mask]     )
   self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.seq_length, self.vocab_size)     )






  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   lowerCamelCase_      =  self.num_labels
   lowerCamelCase_      =  TFEsmForTokenClassification(config=UpperCamelCase     )
   lowerCamelCase_      =  {"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase_      =  model(UpperCamelCase     )
   self.parent.assertEqual(result.logits.shape			,  (self.batch_size, self.seq_length, self.num_labels)     )







  def   snake_case		(    self     ):







   """simple docstring"""
   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      snake_case    (    lowercase			,						lowercase			,						unittest.TestCase			):





  """simple docstring"""




  _lowerCamelCase		    =  (
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  _lowerCamelCase		    =  (
      {
          "feature-extraction": TFEsmModel,
          "fill-mask": TFEsmForMaskedLM,
          "text-classification": TFEsmForSequenceClassification,
          "token-classification": TFEsmForTokenClassification,
          "zero-shot": TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  _lowerCamelCase		    =  False
  _lowerCamelCase		    =  False






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  TFEsmModelTester(self     )
   lowerCamelCase_      =  ConfigTester(self			,  config_class=UpperCamelCase			,  hidden_size=37     )






  def   snake_case		(    self     ):







   """simple docstring"""
   self.config_tester.run_common_tests()






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*UpperCamelCase     )






  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    lowerCamelCase_      =  TFEsmModel.from_pretrained(UpperCamelCase     )
    self.assertIsNotNone(UpperCamelCase     )






  @unittest.skip("Protein models do not support embedding resizing."     )
  def   snake_case		(    self     ):







   """simple docstring"""
   pass






  @unittest.skip("Protein models do not support embedding resizing."     )
  def   snake_case		(    self     ):







   """simple docstring"""
   pass







  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_ ,lowerCamelCase_      =  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_      =  model_class(UpperCamelCase     )
    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(UpperCamelCase			,  UpperCamelCase     )
     for k, v in name.items():
      assert isinstance(UpperCamelCase			,  tf.Variable     )
    else:
     lowerCamelCase_      =  model.get_output_embeddings()
     assert x is None
     lowerCamelCase_      =  model.get_bias()
     assert name is None





@require_tf
class      snake_case    (    unittest.TestCase			):





  """simple docstring"""




  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"     )

   lowerCamelCase_      =  tf.constant([[0, 1, 2, 3, 4, 5]]     )
   lowerCamelCase_      =  model(UpperCamelCase     )[0]
   lowerCamelCase_      =  [1, 6, 33]
   self.assertEqual(list(output.numpy().shape     )			,  UpperCamelCase     )
   # 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   snake_case		(    self     ):







   """simple docstring"""
   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(UpperCamelCase     )[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     )     )



