




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






       def __init__(   self:					str       ,					_lowerCAmelCase:					Dict       ,					_lowerCAmelCase:					Optional[int]=13       ,					_lowerCAmelCase:					int=7       ,					_lowerCAmelCase:					Union[str, Any]=True       ,					_lowerCAmelCase:					List[Any]=True       ,					_lowerCAmelCase:					Optional[Any]=True       ,					_lowerCAmelCase:					Any=True       ,					_lowerCAmelCase:					List[Any]=99       ,					_lowerCAmelCase:					int=32       ,					_lowerCAmelCase:					str=5       ,					_lowerCAmelCase:					str=4       ,					_lowerCAmelCase:					Dict=37       ,					_lowerCAmelCase:					List[str]="gelu"       ,					_lowerCAmelCase:					Optional[int]=0.1       ,					_lowerCAmelCase:					int=0.1       ,					_lowerCAmelCase:					int=5_12       ,					_lowerCAmelCase:					Dict=16       ,					_lowerCAmelCase:					int=2       ,					_lowerCAmelCase:					List[Any]=0.02       ,					_lowerCAmelCase:					Tuple=4       ,					):
         lowercase    :str      	=						parent
         lowercase    :List[str]      	=						batch_size
         lowercase    :str      	=						seq_length
         lowercase    :str      	=						is_training
         lowercase    :Dict      	=						use_attention_mask
         lowercase    :str      	=						use_token_type_ids
         lowercase    :List[str]      	=						use_labels
         lowercase    :Optional[int]      	=						vocab_size
         lowercase    :Tuple      	=						hidden_size
         lowercase    :Optional[int]      	=						num_hidden_layers
         lowercase    :Tuple      	=						num_attention_heads
         lowercase    :Any      	=						intermediate_size
         lowercase    :str      	=						hidden_act
         lowercase    :Any      	=						hidden_dropout_prob
         lowercase    :Optional[Any]      	=						attention_probs_dropout_prob
         lowercase    :Optional[int]      	=						max_position_embeddings
         lowercase    :Tuple      	=						type_vocab_size
         lowercase    :List[str]      	=						type_sequence_label_size
         lowercase    :List[Any]      	=						initializer_range
         lowercase    :str      	=						num_choices






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :Optional[Any]      	=						ids_tensor([self.batch_size, self.seq_length]       ,					self.vocab_size				)

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

         lowercase    :Dict      	=						None
         if self.use_token_type_ids:
           lowercase    :List[str]      	=						ids_tensor([self.batch_size, self.seq_length]       ,					self.type_vocab_size				)

         lowercase    :List[Any]      	=						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=lowerCAmelCase__       ,					initializer_range=self.initializer_range       ,					)

         return config, input_ids, token_type_ids, attention_mask






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Union[str, Any]      	=						self.prepare_config_and_inputs()
         lowercase    :int      	=						config_and_inputs
         lowercase    :Optional[Any]      	=						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
         return config, inputs_dict







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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Dict      	=						FlaxAlbertModelTester(self				)






       @slow
       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         for model_class_name in self.all_model_classes:
           lowercase    :Optional[int]      	=						model_class_name.from_pretrained("albert-base-v2"				)
           lowercase    :Tuple      	=						model(np.ones((1, 1)				)				)
           self.assertIsNotNone(lowerCAmelCase__				)







@require_flax
class    __lowerCAmelCase       (							unittest.TestCase):






       @slow
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Any      	=						FlaxAlbertModel.from_pretrained("albert-base-v2"				)
         lowercase    :Union[str, Any]      	=						np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]				)
         lowercase    :int      	=						np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
         lowercase    :Any      	=						model(lowerCAmelCase__       ,					attention_mask=lowerCAmelCase__				)[0]
         lowercase    :List[Any]      	=						(1, 11, 7_68)
         self.assertEqual(output.shape       ,					lowerCAmelCase__				)
         lowercase    :Optional[int]      	=						np.array(
             [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]				)

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









from collections.abc import Generator
from math import sin



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if len(lowerCamelCase   ) != 32:
    raise ValueError("Input must be of length 32"   )

  lowercase    :Any      	=						B""
  for i in [3, 2, 1, 0]:
    little_endian += string_aa[8 * i : 8 * i + 8]
  return little_endian



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )

  lowercase    :List[Any]      	=						format(lowerCamelCase,			"08x"   )[-8:]
  lowercase    :Tuple      	=						B""
  for i in [3, 2, 1, 0]:
    little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("utf-8"   )
  return little_endian_hex



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Union[str, Any]      	=						B""
  for char in message:
    bit_string += format(lowerCamelCase,			"08b"   ).encode("utf-8"   )
  lowercase    :str      	=						format(len(lowerCamelCase   ),			"064b"   ).encode("utf-8"   )

  # Pad bit_string to a multiple of 512 chars
  bit_string += b"1"
  while len(lowerCamelCase   ) % 512 != 448:
    bit_string += b"0"
  bit_string += to_little_endian(start_len[32:]   ) + to_little_endian(start_len[:32]   )

  return bit_string



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if len(lowerCamelCase   ) % 512 != 0:
    raise ValueError("Input must have length that's a multiple of 512"   )

  for pos in range(0,			len(lowerCamelCase   ),			512   ):
    lowercase    :List[str]      	=						bit_string[pos : pos + 512]
    lowercase    :int      	=						[]
    for i in range(0,			512,			32   ):
      block_words.append(int(to_little_endian(block[i : i + 32]   ),			2   )   )
    yield block_words



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )

  lowercase    :List[str]      	=						format(lowerCamelCase,			"032b"   )
  lowercase    :Optional[Any]      	=						""
  for c in i_str:
    new_str += "1" if c == "0" else "0"
  return int(lowerCamelCase,			2   )



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  return (a + b) % 2**32



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )
  if shift < 0:
    raise ValueError("Shift must be non-negative"   )
  return ((i << shift) ^ (i >> (32 - shift))) % 2**32



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Tuple      	=						preprocess(lowerCamelCase   )

  lowercase    :Tuple      	=						[int(2**32 * abs(sin(i + 1   )   )   ) for i in range(64   )]

  # Starting states
  lowercase    :Optional[Any]      	=						0X6745_2301
  lowercase    :Union[str, Any]      	=						0XEFCD_AB89
  lowercase    :Optional[int]      	=						0X98BA_DCFE
  lowercase    :str      	=						0X1032_5476

  lowercase    :str      	=						[
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
  ]

  # Process bit string in chunks, each with 16 32-char words
  for block_words in get_block_words(lowerCamelCase   ):
    lowercase    :str      	=						aa
    lowercase    :int      	=						ba
    lowercase    :int      	=						ca
    lowercase    :List[str]      	=						da

    # Hash current chunk
    for i in range(64   ):
      if i <= 15:
        # f = (b & c) | (not_32(b) & d)     # Alternate definition for f
        lowercase    :Optional[int]      	=						d ^ (b & (c ^ d))
        lowercase    :Optional[Any]      	=						i
      elif i <= 31:
        # f = (d & b) | (not_32(d) & c)     # Alternate definition for f
        lowercase    :Any      	=						c ^ (d & (b ^ c))
        lowercase    :Union[str, Any]      	=						(5 * i + 1) % 16
      elif i <= 47:
        lowercase    :Any      	=						b ^ c ^ d
        lowercase    :List[Any]      	=						(3 * i + 5) % 16
      else:
        lowercase    :List[Any]      	=						c ^ (b | not_aa(lowerCamelCase   ))
        lowercase    :Any      	=						(7 * i) % 16
      lowercase    :Dict      	=						(f + a + added_consts[i] + block_words[g]) % 2**32
      lowercase    :Any      	=						d
      lowercase    :str      	=						c
      lowercase    :Union[str, Any]      	=						b
      lowercase    :Union[str, Any]      	=						sum_aa(lowerCamelCase,			left_rotate_aa(lowerCamelCase,			shift_amounts[i]   )   )

    # Add hashed chunk to running total
    lowercase    :Optional[Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :List[Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :List[str]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :Union[str, Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )

  lowercase    :int      	=						reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   )
  return digest


if __name__ == "__main__":
     import doctest

     doctest.testmod()



