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   lowercase_    :



	"""simple docstring"""





	def __init__(				self						,     __SCREAMING_SNAKE_CASE						,     )		->int:
								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.0_2
								lowerCAmelCase      						=						3
								lowerCAmelCase      						=						4
								lowerCAmelCase      						=						None


	def      SCREAMING_SNAKE_CASE_							(				self							)		->Optional[int]:
								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							)		->List[str]:
								(
								    (
								    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						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE							)		->Optional[int]:
								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						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     )		->Dict:
								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						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE							)		->int:
								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						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE						,     __SCREAMING_SNAKE_CASE							)		->Tuple:
								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							)		->Union[str, 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   lowercase_    (    __snake_case ,	__snake_case ,	unittest.TestCase		):



	"""simple docstring"""





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


	def      SCREAMING_SNAKE_CASE_							(				self							)		->Any:
								lowerCAmelCase      						=						TFEsmModelTester(self							)
								lowerCAmelCase      						=						ConfigTester(self						,     config_class=__lowerCamelCase						,     hidden_size=37							)


	def      SCREAMING_SNAKE_CASE_							(				self							)		->Optional[int]:
								self.config_tester.run_common_tests()


	def      SCREAMING_SNAKE_CASE_							(				self							)		->Optional[Any]:
								lowerCAmelCase      						=						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__lowerCamelCase							)


	def      SCREAMING_SNAKE_CASE_							(				self							)		->Optional[Any]:
								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							)		->List[str]:
								lowerCAmelCase      						=						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase							)


	@slow
	def      SCREAMING_SNAKE_CASE_							(				self							)		->Optional[Any]:







								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							)		->Optional[int]:
								pass


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




	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(__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   lowercase_    (    unittest.TestCase		):



	"""simple docstring"""





	@slow
	def      SCREAMING_SNAKE_CASE_							(				self							)		->int:
								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.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      SCREAMING_SNAKE_CASE_							(				self							)		->List[str]:
								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.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							)							)







import argparse
import json
from typing import List

from ltp import LTP

from transformers.models.bert.tokenization_bert import BertTokenizer
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->Optional[Any]:
   # This defines a "chinese character" as anything in the CJK Unicode block:
   #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
   #
   # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
   # despite its name. The modern Korean Hangul alphabet is a different block,
   # as is Japanese Hiragana and Katakana. Those alphabets are used to write
   # space-separated words, so they are not treated specially and handled
   # like the all of the other languages.
   if (
       (cp >= 0x4_E_0_0 and cp <= 0x9_F_F_F)
       or (cp >= 0x3_4_0_0 and cp <= 0x4_D_B_F)  #
       or (cp >= 0x2_0_0_0_0 and cp <= 0x2_A_6_D_F)  #
       or (cp >= 0x2_A_7_0_0 and cp <= 0x2_B_7_3_F)  #
       or (cp >= 0x2_B_7_4_0 and cp <= 0x2_B_8_1_F)  #
       or (cp >= 0x2_B_8_2_0 and cp <= 0x2_C_E_A_F)  #
       or (cp >= 0xF_9_0_0 and cp <= 0xF_A_F_F)
       or (cp >= 0x2_F_8_0_0 and cp <= 0x2_F_A_1_F)  #
   ):  #
      return True

   return False
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->List[str]:
   # word like '180' or '身高' or '神'
   for char in word:
      UpperCAmelCase      					=  ord(lowerCAmelCase_ )
      if not _is_chinese_char(lowerCAmelCase_ ):
         return 0
   return 1
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->Optional[int]:
   UpperCAmelCase      					=  set()

   for token in tokens:
      UpperCAmelCase      					=  len(lowerCAmelCase_ ) > 1 and is_chinese(lowerCAmelCase_ )
      if chinese_word:
         word_set.add(lowerCAmelCase_ )
   UpperCAmelCase      					=  list(lowerCAmelCase_ )
   return word_list
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Union[str, Any]:
   if not chinese_word_set:
      return bert_tokens
   UpperCAmelCase      					=  max([len(lowerCAmelCase_ ) for w in chinese_word_set] )

   UpperCAmelCase      					=  bert_tokens
   UpperCAmelCase					,	UpperCAmelCase      					=  0, len(lowerCAmelCase_ )
   while start < end:
      UpperCAmelCase      					=  True
      if is_chinese(bert_word[start] ):
         UpperCAmelCase      					=  min(end - start	,    lowerCAmelCase_ )
         for i in range(lowerCAmelCase_	,    1	,    -1 ):
            UpperCAmelCase      					=  """""".join(bert_word[start : start + i] )
            if whole_word in chinese_word_set:
               for j in range(start + 1	,    start + i ):
                  UpperCAmelCase      					=  """##""" + bert_word[j]
               UpperCAmelCase      					=  start + i
               UpperCAmelCase      					=  False
               break
      if single_word:
         start += 1
   return bert_word
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Dict:
   UpperCAmelCase      					=  []

   for i in range(0	,    len(lowerCAmelCase_ )	,    1_0_0 ):
      UpperCAmelCase      					=  ltp_tokenizer.pipeline(lines[i : i + 1_0_0]	,    tasks=["""cws"""] ).cws
      UpperCAmelCase      					=  [get_chinese_word(lowerCAmelCase_ ) for r in res]
      ltp_res.extend(lowerCAmelCase_ )
   assert len(lowerCAmelCase_ ) == len(lowerCAmelCase_ )

   UpperCAmelCase      					=  []
   for i in range(0	,    len(lowerCAmelCase_ )	,    1_0_0 ):
      UpperCAmelCase      					=  bert_tokenizer(lines[i : i + 1_0_0]	,    add_special_tokens=lowerCAmelCase_	,    truncation=lowerCAmelCase_	,    max_length=5_1_2 )
      bert_res.extend(res["""input_ids"""] )
   assert len(lowerCAmelCase_ ) == len(lowerCAmelCase_ )

   UpperCAmelCase      					=  []
   for input_ids, chinese_word in zip(lowerCAmelCase_	,    lowerCAmelCase_ ):
      UpperCAmelCase      					=  []
      for id in input_ids:
         UpperCAmelCase      					=  bert_tokenizer._convert_id_to_token(lowerCAmelCase_ )
         input_tokens.append(lowerCAmelCase_ )
      UpperCAmelCase      					=  add_sub_symbol(lowerCAmelCase_	,    lowerCAmelCase_ )
      UpperCAmelCase      					=  []
      # We only save pos of chinese subwords start with ##, which mean is part of a whole word.
      for i, token in enumerate(lowerCAmelCase_ ):
         if token[:2] == "##":
            UpperCAmelCase      					=  token[2:]
            # save chinese tokens' pos
            if len(lowerCAmelCase_ ) == 1 and _is_chinese_char(ord(lowerCAmelCase_ ) ):
               ref_id.append(lowerCAmelCase_ )
      ref_ids.append(lowerCAmelCase_ )

   assert len(lowerCAmelCase_ ) == len(lowerCAmelCase_ )

   return ref_ids
def 	_UpperCamelCase					(    lowerCAmelCase_ )    ->int:
   # For Chinese (Ro)Bert, the best result is from : RoBERTa-wwm-ext (https://github.com/ymcui/Chinese-BERT-wwm)
   # If we want to fine-tune these model, we have to use same tokenizer : LTP (https://github.com/HIT-SCIR/ltp)
   with open(args.file_name	,    """r"""	,    encoding="""utf-8""" ) as f:
      UpperCAmelCase      					=  f.readlines()
   UpperCAmelCase      					=  [line.strip() for line in data if len(lowerCAmelCase_ ) > 0 and not line.isspace()]  # avoid delimiter like '\u2029'
   UpperCAmelCase      					=  LTP(args.ltp )  # faster in GPU device
   UpperCAmelCase      					=  BertTokenizer.from_pretrained(args.bert )

   UpperCAmelCase      					=  prepare_ref(lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )

   with open(args.save_path	,    """w"""	,    encoding="""utf-8""" ) as f:
      UpperCAmelCase      					=  [json.dumps(lowerCAmelCase_ ) + """\n""" for ref in ref_ids]
      f.writelines(lowerCAmelCase_ )


if __name__ == "__main__":
 __a	 =      argparse.ArgumentParser(description="""prepare_chinese_ref""")
 parser.add_argument(
     """--file_name""",
     required=False,
     type=str,
     default="""./resources/chinese-demo.txt""",
     help="""file need process, same as training data in lm""",
 )
 parser.add_argument(
     """--ltp""",
     required=False,
     type=str,
     default="""./resources/ltp""",
     help="""resources for LTP tokenizer, usually a path""",
 )
 parser.add_argument(
     """--bert""",
     required=False,
     type=str,
     default="""./resources/robert""",
     help="""resources for Bert tokenizer""",
 )
 parser.add_argument(
     """--save_path""",
     required=False,
     type=str,
     default="""./resources/ref.txt""",
     help="""path to save res""",
 )

 __a	 =      parser.parse_args()
 main(args)



