


"""simple docstring"""





import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax.numpy as jnp

 from transformers.models.albert.modeling_flax_albert import (
     FlaxAlbertForMaskedLM,
     FlaxAlbertForMultipleChoice,
     FlaxAlbertForPreTraining,
     FlaxAlbertForQuestionAnswering,
     FlaxAlbertForSequenceClassification,
     FlaxAlbertForTokenClassification,
     FlaxAlbertModel,
 )



class 						snake_case_(				unittest.TestCase	):






     def __init__(		self     :    Tuple    , UpperCamelCase_     :    int    , UpperCamelCase_     :    List[Any]=1_3    , UpperCamelCase_     :    Tuple=7    , UpperCamelCase_     :    int=True    , UpperCamelCase_     :    Dict=True    , UpperCamelCase_     :    List[Any]=True    , UpperCamelCase_     :    Optional[int]=True    , UpperCamelCase_     :    List[str]=9_9    , UpperCamelCase_     :    Optional[int]=3_2    , UpperCamelCase_     :    str=5    , UpperCamelCase_     :    Optional[Any]=4    , UpperCamelCase_     :    List[str]=3_7    , UpperCamelCase_     :    List[Any]="gelu"    , UpperCamelCase_     :    Optional[int]=0.1    , UpperCamelCase_     :    Optional[int]=0.1    , UpperCamelCase_     :    Optional[Any]=5_1_2    , UpperCamelCase_     :    Optional[int]=1_6    , UpperCamelCase_     :    Dict=2    , UpperCamelCase_     :    str=0.02    , UpperCamelCase_     :    Any=4    , ):
           lowerCAmelCase    :		List[str]				=      parent
           lowerCAmelCase    :		List[Any]				=      batch_size
           lowerCAmelCase    :		Any				=      seq_length
           lowerCAmelCase    :		Dict				=      is_training
           lowerCAmelCase    :		Any				=      use_attention_mask
           lowerCAmelCase    :		Tuple				=      use_token_type_ids
           lowerCAmelCase    :		Any				=      use_labels
           lowerCAmelCase    :		Any				=      vocab_size
           lowerCAmelCase    :		Any				=      hidden_size
           lowerCAmelCase    :		Any				=      num_hidden_layers
           lowerCAmelCase    :		str				=      num_attention_heads
           lowerCAmelCase    :		Optional[int]				=      intermediate_size
           lowerCAmelCase    :		Tuple				=      hidden_act
           lowerCAmelCase    :		int				=      hidden_dropout_prob
           lowerCAmelCase    :		Optional[int]				=      attention_probs_dropout_prob
           lowerCAmelCase    :		Union[str, Any]				=      max_position_embeddings
           lowerCAmelCase    :		Any				=      type_vocab_size
           lowerCAmelCase    :		str				=      type_sequence_label_size
           lowerCAmelCase    :		List[str]				=      initializer_range
           lowerCAmelCase    :		Union[str, Any]				=      num_choices






     def    lowerCamelCase__   (		self     :    Tuple ):
           lowerCAmelCase    :		List[Any]				=      ids_tensor([self.batch_size, self.seq_length]    , self.vocab_size )

           lowerCAmelCase    :		int				=      None
           if self.use_attention_mask:
                 lowerCAmelCase    :		List[Any]				=      random_attention_mask([self.batch_size, self.seq_length] )

           lowerCAmelCase    :		Union[str, Any]				=      None
           if self.use_token_type_ids:
                 lowerCAmelCase    :		Dict				=      ids_tensor([self.batch_size, self.seq_length]    , self.type_vocab_size )

           lowerCAmelCase    :		Optional[int]				=      AlbertConfig(
               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    , is_decoder=UpperCamelCase_    , initializer_range=self.initializer_range    , )

           return config, input_ids, token_type_ids, attention_mask






     def    lowerCamelCase__   (		self     :    str ):
           lowerCAmelCase    :		Any				=      self.prepare_config_and_inputs()
           lowerCAmelCase,					lowerCAmelCase,					lowerCAmelCase,					lowerCAmelCase    :		int				=      config_and_inputs
           lowerCAmelCase    :		Optional[Any]				=      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
           return config, inputs_dict



@require_flax
class 						snake_case_(				a__							,    unittest.TestCase	):
     __UpperCamelCase     =    (
         (
             FlaxAlbertModel,
             FlaxAlbertForPreTraining,
             FlaxAlbertForMaskedLM,
             FlaxAlbertForMultipleChoice,
             FlaxAlbertForQuestionAnswering,
             FlaxAlbertForSequenceClassification,
             FlaxAlbertForTokenClassification,
             FlaxAlbertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )






     def    lowerCamelCase__   (		self     :    int ):
           lowerCAmelCase    :		Optional[Any]				=      FlaxAlbertModelTester(self )






     @slow
     def    lowerCamelCase__   (		self     :    str ):
           for model_class_name in self.all_model_classes:
                 lowerCAmelCase    :		Union[str, Any]				=      model_class_name.from_pretrained('''albert-base-v2''' )
                 lowerCAmelCase    :		List[str]				=      model(np.ones((1, 1) ) )
                 self.assertIsNotNone(UpperCamelCase_ )



@require_flax
class 						snake_case_(				unittest.TestCase	):






     @slow
     def    lowerCamelCase__   (		self     :    int ):
           lowerCAmelCase    :		Tuple				=      FlaxAlbertModel.from_pretrained('''albert-base-v2''' )
           lowerCAmelCase    :		List[str]				=      np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]] )
           lowerCAmelCase    :		List[Any]				=      np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] )
           lowerCAmelCase    :		int				=      model(UpperCamelCase_    , attention_mask=UpperCamelCase_ )[0]
           lowerCAmelCase    :		Optional[Any]				=      (1, 1_1, 7_6_8)
           self.assertEqual(output.shape    , UpperCamelCase_ )
           lowerCAmelCase    :		Dict				=      np.array(
               [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]] )

           self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]    , UpperCamelCase_    , atol=1E-4 ) )






"""simple docstring"""





import math


def    _snake_case  (  ):
      lowerCAmelCase    :		Union[str, Any]				=      input('''Enter message: '''	)
      lowerCAmelCase    :		Optional[int]				=      int(input(f'''Enter key [2-{len(_snake_case	) - 1}]: '''	)	)
      lowerCAmelCase    :		str				=      input('''Encryption/Decryption [e/d]: '''	)

      if mode.lower().startswith('''e'''	):
            lowerCAmelCase    :		Any				=      encrypt_message(_snake_case       ,      _snake_case	)
      elif mode.lower().startswith('''d'''	):
            lowerCAmelCase    :		Union[str, Any]				=      decrypt_message(_snake_case       ,      _snake_case	)

      # Append pipe symbol (vertical bar) to identify spaces at the end.
      print(f'''Output:\n{text + "|"}'''	)


def    _snake_case  (  _snake_case :	int       ,      _snake_case :	str	):
      lowerCAmelCase    :		Optional[Any]				=      [''''''] * key
      for col in range(_snake_case	):
            lowerCAmelCase    :		Optional[Any]				=      col
            while pointer < len(_snake_case	):
                  cipher_text[col] += message[pointer]
                  pointer += key
      return "".join(_snake_case	)


def    _snake_case  (  _snake_case :	int       ,      _snake_case :	str	):
      lowerCAmelCase    :		Union[str, Any]				=      math.ceil(len(_snake_case	) / key	)
      lowerCAmelCase    :		str				=      key
      lowerCAmelCase    :		Any				=      (num_cols * num_rows) - len(_snake_case	)
      lowerCAmelCase    :		Dict				=      [''''''] * num_cols
      lowerCAmelCase    :		int				=      0
      lowerCAmelCase    :		int				=      0

      for symbol in message:
            plain_text[col] += symbol
            col += 1

            if (
                (col == num_cols)
                or (col == num_cols - 1)
                and (row >= num_rows - num_shaded_boxes)
            ):
                  lowerCAmelCase    :		int				=      0
                  row += 1

      return "".join(_snake_case	)


if __name__ == "__main__":
 import doctest

 doctest.testmod()
 main()


