


import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      ErnieMForInformationExtraction,
      ErnieMForMultipleChoice,
      ErnieMForQuestionAnswering,
      ErnieMForSequenceClassification,
      ErnieMForTokenClassification,
      ErnieMModel,
  )
  from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST

class      lowerCamelCase	:

       def __init__(			self	,						a__	,						a__=13	,						a__=7	,						a__=True	,						a__=True	,						a__=True	,						a__=99	,						a__=32	,						a__=5	,						a__=4	,						a__=37	,						a__="gelu"	,						a__=0.1	,						a__=0.1	,						a__=5_12	,						a__=16	,						a__=2	,						a__=0.02	,						a__=3	,						a__=4	,						a__=None	,						):






        '''simple docstring'''






        snake_case   : int	  =      parent
        snake_case   : List[Any]	  =      batch_size
        snake_case   : int	  =      seq_length
        snake_case   : Optional[int]	  =      is_training
        snake_case   : List[str]	  =      use_input_mask
        snake_case   : Tuple	  =      use_labels
        snake_case   : List[Any]	  =      vocab_size
        snake_case   : Optional[int]	  =      hidden_size
        snake_case   : Optional[int]	  =      num_hidden_layers
        snake_case   : Optional[int]	  =      num_attention_heads
        snake_case   : Tuple	  =      intermediate_size
        snake_case   : Tuple	  =      hidden_act
        snake_case   : Optional[Any]	  =      hidden_dropout_prob
        snake_case   : Dict	  =      attention_probs_dropout_prob
        snake_case   : List[Any]	  =      max_position_embeddings
        snake_case   : Dict	  =      type_vocab_size
        snake_case   : Tuple	  =      type_sequence_label_size
        snake_case   : Any	  =      initializer_range
        snake_case   : Tuple	  =      num_labels
        snake_case   : List[Any]	  =      num_choices
        snake_case   : int	  =      scope

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Union[str, Any]	  =      ids_tensor([self.batch_size, self.seq_length]	,						self.vocab_size					)

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

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

        snake_case   : Tuple	  =      self.get_config()

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

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Tuple	  =      ids_tensor([self.batch_size, self.seq_length]	,						self.vocab_size					)

        snake_case   : int	  =      None
        if self.use_input_mask:
         snake_case   : Union[str, Any]	  =      random_attention_mask([self.batch_size, self.seq_length]					)
        snake_case   : Dict	  =      self.get_config()

        return config, input_ids, input_mask

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        return ErnieMConfig(
            vocab_size=self.vocab_size	,						hidden_size=self.hidden_size	,						num_hidden_layers=self.num_hidden_layers	,						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	,						)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : Optional[int]	  =      ErnieMModel(config=a__					)
        model.to(a__					)
        model.eval()
        snake_case   : Optional[int]	  =      model(a__	,						return_dict=a__					)
        self.parent.assertEqual(result.last_hidden_state.shape	,						(self.batch_size, self.seq_length, self.hidden_size)					)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : Any	  =      ErnieMForQuestionAnswering(config=a__					)
        model.to(a__					)
        model.eval()
        snake_case   : Dict	  =      model(
            a__	,						attention_mask=a__	,						start_positions=a__	,						end_positions=a__	,						)
        self.parent.assertEqual(result.start_logits.shape	,						(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape	,						(self.batch_size, self.seq_length)					)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : List[Any]	  =      ErnieMForInformationExtraction(config=a__					)
        model.to(a__					)
        model.eval()
        snake_case   : Any	  =      torch.ones_like(a__	,						dtype=torch.floataa					)
        snake_case   : Dict	  =      model(
            a__	,						attention_mask=a__	,						start_positions=a__	,						end_positions=a__	,						)
        self.parent.assertEqual(result.start_logits.shape	,						(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape	,						(self.batch_size, self.seq_length)					)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : List[str]	  =      self.num_labels
        snake_case   : Optional[int]	  =      ErnieMForSequenceClassification(a__					)
        model.to(a__					)
        model.eval()
        snake_case   : str	  =      model(a__	,						attention_mask=a__	,						labels=a__					)
        self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.num_labels)					)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : Any	  =      self.num_labels
        snake_case   : Optional[Any]	  =      ErnieMForTokenClassification(config=a__					)
        model.to(a__					)
        model.eval()
        input_ids.to(a__					)
        input_mask.to(a__					)
        token_labels.to(a__					)

        snake_case   : int	  =      model(a__	,						attention_mask=a__	,						labels=a__					)

        self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.seq_length, self.num_labels)					)

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        snake_case   : Tuple	  =      self.num_choices
        snake_case   : List[str]	  =      ErnieMForMultipleChoice(config=a__					)
        model.to(a__					)
        model.eval()
        snake_case   : str	  =      input_ids.unsqueeze(1					).expand(-1	,						self.num_choices	,						-1					).contiguous()
        snake_case   : List[Any]	  =      input_mask.unsqueeze(1					).expand(-1	,						self.num_choices	,						-1					).contiguous()
        snake_case   : str	  =      model(
            a__	,						attention_mask=a__	,						labels=a__	,						)
        self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.num_choices)					)



       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Optional[int]	  =      self.prepare_config_and_inputs()
        (
            (
            snake_case
        )  ,					(
            snake_case
        )  ,					(
            snake_case
        )  ,					(
            snake_case
        )  ,					(
            snake_case
        )  ,					(
            snake_case
        )  ,					
        )   : int	  =      config_and_inputs
        snake_case   : Any	  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict







@require_torch
class      lowerCamelCase	(  SCREAMING_SNAKE_CASE,			SCREAMING_SNAKE_CASE,			unittest.TestCase     ):
       _UpperCamelCase    		=      (
           (
               ErnieMModel,
               ErnieMForMultipleChoice,
               ErnieMForQuestionAnswering,
               ErnieMForSequenceClassification,
               ErnieMForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       _UpperCamelCase    		=      ()
       _UpperCamelCase    		=      (
           {
               """feature-extraction""": ErnieMModel,
               """question-answering""": ErnieMForQuestionAnswering,
               """text-classification""": ErnieMForSequenceClassification,
               """token-classification""": ErnieMForTokenClassification,
               """zero-shot""": ErnieMForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       _UpperCamelCase    		=      False

       def 				UpperCAmelCase				(			self	,						a__	,						a__	,						a__	,						a__	,						a__					):






        '''simple docstring'''






        if pipeline_test_casse_name == "QAPipelineTests":
         return True

        return False

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : int	  =      ErnieMModelTester(self					)
        snake_case   : int	  =      ConfigTester(self	,						config_class=a__	,						hidden_size=37					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        self.config_tester.run_common_tests()

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Union[str, Any]	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Union[str, Any]	  =      self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
         snake_case   : List[Any]	  =      type
         self.model_tester.create_and_check_model(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Dict	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : List[Any]	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : int	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Dict	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_information_extraction(*a__					)

       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Optional[Any]	  =      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*a__					)



       @slow
       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         snake_case   : List[Any]	  =      ErnieMModel.from_pretrained(a__					)
         self.assertIsNotNone(a__					)







@require_torch
class      lowerCamelCase	(  unittest.TestCase     ):

       @slow
       def 				UpperCAmelCase				(			self					):






        '''simple docstring'''






        snake_case   : Union[str, Any]	  =      ErnieMModel.from_pretrained('''susnato/ernie-m-base_pytorch'''					)
        model.eval()
        snake_case   : Dict	  =      torch.tensor([[0, 1, 2, 3, 4, 5]]					)
        snake_case   : List[Any]	  =      model(a__					)[0]

        # TODO Replace vocab size
        snake_case   : Union[str, Any]	  =      7_68

        snake_case   : Optional[Any]	  =      torch.Size((1, 6, hidden_size)					)
        self.assertEqual(output.shape	,						a__					)

        snake_case   : Any	  =      torch.tensor(
            [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]					)

        self.assertTrue(torch.allclose(output[:, :3, :3]	,						a__	,						atol=1e-3					)					)










import os
import sys

from . import rsa_key_generator as rkg

a				:							str					=			128
a				:							str					=			256






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase = DEFAULT_BLOCK_SIZE       )		->list[int]:


 '''simple docstring'''
 snake_case   : List[str]	  =      message.encode('''ascii'''       )

 snake_case   : List[Any]	  =      []
 for block_start in range(0 ,     len(__lowerCamelCase       ) ,     __lowerCamelCase       ):
  snake_case   : int	  =      0
  for i in range(__lowerCamelCase ,     min(block_start + block_size ,     len(__lowerCamelCase       )       )       ):
   block_int += message_bytes[i] * (BYTE_SIZE ** (i % block_size))
  block_ints.append(__lowerCamelCase       )
 return block_ints






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase = DEFAULT_BLOCK_SIZE       )		->str:


 '''simple docstring'''
 snake_case   : list[str]	  =      []
 for block_int in block_ints:
  snake_case   : list[str]	  =      []
  for i in range(block_size - 1 ,     -1 ,     -1       ):
   if len(__lowerCamelCase       ) + i < message_length:
    snake_case   : Any	  =      block_int // (BYTE_SIZE**i)
    snake_case   : Optional[Any]	  =      block_int % (BYTE_SIZE**i)
    block_message.insert(0 ,     chr(__lowerCamelCase       )       )
  message.extend(__lowerCamelCase       )
 return "".join(__lowerCamelCase       )






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase = DEFAULT_BLOCK_SIZE       )		->list[int]:


 '''simple docstring'''
 snake_case   : str	  =      []
 snake_case  ,					snake_case   : Optional[Any]	  =      key

 for block in get_blocks_from_text(__lowerCamelCase ,     __lowerCamelCase       ):
  encrypted_blocks.append(pow(__lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase       )       )
 return encrypted_blocks






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase = DEFAULT_BLOCK_SIZE ,     )		->str:


 '''simple docstring'''
 snake_case   : List[Any]	  =      []
 snake_case  ,					snake_case   : List[str]	  =      key
 for block in encrypted_blocks:
  decrypted_blocks.append(pow(__lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase       )       )
 return get_text_from_blocks(__lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase       )






def    _UpperCAmelCase	(   __lowerCamelCase       )		->tuple[int, int, int]:


 '''simple docstring'''
 with open(__lowerCamelCase       ) as fo:
  snake_case   : List[Any]	  =      fo.read()
 snake_case  ,					snake_case  ,					snake_case   : Dict	  =      content.split(''','''       )
 return (int(__lowerCamelCase       ), int(__lowerCamelCase       ), int(__lowerCamelCase       ))






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase = DEFAULT_BLOCK_SIZE ,     )		->str:


 '''simple docstring'''
 snake_case  ,					snake_case  ,					snake_case   : List[Any]	  =      read_key_file(__lowerCamelCase       )
 if key_size < block_size * 8:
  sys.exit(
      '''ERROR: Block size is {} bits and key size is {} bits. The RSA cipher '''
      '''requires the block size to be equal to or greater than the key size. '''
      '''Either decrease the block size or use different keys.'''.format(
          block_size * 8 ,     __lowerCamelCase       )       )

 snake_case   : Any	  =      [str(__lowerCamelCase       ) for i in encrypt_message(__lowerCamelCase ,     (n, e) ,     __lowerCamelCase       )]

 snake_case   : List[str]	  =      ''','''.join(__lowerCamelCase       )
 snake_case   : str	  =      f'''{len(__lowerCamelCase       )}_{block_size}_{encrypted_content}'''
 with open(__lowerCamelCase ,     '''w'''       ) as fo:
  fo.write(__lowerCamelCase       )
 return encrypted_content






def    _UpperCAmelCase	(   __lowerCamelCase ,     __lowerCamelCase       )		->str:


 '''simple docstring'''
 snake_case  ,					snake_case  ,					snake_case   : str	  =      read_key_file(__lowerCamelCase       )
 with open(__lowerCamelCase       ) as fo:
  snake_case   : Dict	  =      fo.read()
 snake_case  ,					snake_case  ,					snake_case   : Union[str, Any]	  =      content.split('''_'''       )
 snake_case   : str	  =      int(__lowerCamelCase       )
 snake_case   : int	  =      int(__lowerCamelCase       )

 if key_size < block_size * 8:
  sys.exit(
      '''ERROR: Block size is {} bits and key size is {} bits. The RSA cipher '''
      '''requires the block size to be equal to or greater than the key size. '''
      '''Did you specify the correct key file and encrypted file?'''.format(
          block_size * 8 ,     __lowerCamelCase       )       )

 snake_case   : str	  =      []
 for block in encrypted_message.split(''','''       ):
  encrypted_blocks.append(int(__lowerCamelCase       )       )

 return decrypt_message(__lowerCamelCase ,     __lowerCamelCase ,     (n, d) ,     __lowerCamelCase       )






def    _UpperCAmelCase	(   )		->None:


 '''simple docstring'''
 snake_case   : List[Any]	  =      '''encrypted_file.txt'''
 snake_case   : Optional[int]	  =      input(R'''Encrypt\Decrypt [e\d]: '''       )

 if response.lower().startswith('''e'''       ):
  snake_case   : Dict	  =      '''encrypt'''
 elif response.lower().startswith('''d'''       ):
  snake_case   : Dict	  =      '''decrypt'''

 if mode == "encrypt":
  if not os.path.exists('''rsa_pubkey.txt'''       ):
   rkg.make_key_files('''rsa''' ,     1024       )

  snake_case   : List[Any]	  =      input('''\nEnter message: '''       )
  snake_case   : Union[str, Any]	  =      '''rsa_pubkey.txt'''
  print(f'''Encrypting and writing to {filename}...'''       )
  snake_case   : List[str]	  =      encrypt_and_write_to_file(__lowerCamelCase ,     __lowerCamelCase ,     __lowerCamelCase       )

  print('''\nEncrypted text:'''       )
  print(__lowerCamelCase       )

 elif mode == "decrypt":
  snake_case   : str	  =      '''rsa_privkey.txt'''
  print(f'''Reading from {filename} and decrypting...'''       )
  snake_case   : int	  =      read_from_file_and_decrypt(__lowerCamelCase ,     __lowerCamelCase       )
  print('''writing decryption to rsa_decryption.txt...'''       )
  with open('''rsa_decryption.txt''' ,     '''w'''       ) as dec:
   dec.write(__lowerCamelCase       )

  print('''\nDecryption:'''       )
  print(__lowerCamelCase       )


if __name__ == "__main__":
  main()






