


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





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=13    ,						UpperCamelCase=7    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=99    ,						UpperCamelCase=32    ,						UpperCamelCase=5    ,						UpperCamelCase=4    ,						UpperCamelCase=37    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=512    ,						UpperCamelCase=16    ,						UpperCamelCase=2    ,						UpperCamelCase=0.02    ,						UpperCamelCase=4    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						seq_length
   lowerCamelCase_  				=						is_training
   lowerCamelCase_  				=						use_attention_mask
   lowerCamelCase_  				=						use_token_type_ids
   lowerCamelCase_  				=						use_labels
   lowerCamelCase_  				=						vocab_size
   lowerCamelCase_  				=						hidden_size
   lowerCamelCase_  				=						num_hidden_layers
   lowerCamelCase_  				=						num_attention_heads
   lowerCamelCase_  				=						intermediate_size
   lowerCamelCase_  				=						hidden_act
   lowerCamelCase_  				=						hidden_dropout_prob
   lowerCamelCase_  				=						attention_probs_dropout_prob
   lowerCamelCase_  				=						max_position_embeddings
   lowerCamelCase_  				=						type_vocab_size
   lowerCamelCase_  				=						type_sequence_label_size
   lowerCamelCase_  				=						initializer_range
   lowerCamelCase_  				=						num_choices






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size  )

   lowerCamelCase_  				=						None
   if self.use_attention_mask:
    lowerCamelCase_  				=						random_attention_mask([self.batch_size, self.seq_length]  )

   lowerCamelCase_  				=						None
   if self.use_token_type_ids:
    lowerCamelCase_  				=						ids_tensor([self.batch_size, self.seq_length]    ,						self.type_vocab_size  )

   lowerCamelCase_  				=						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 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.prepare_config_and_inputs()
   lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_  				=						config_and_inputs
   lowerCamelCase_  				=						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
   return config, inputs_dict





@require_flax
class   snake_case					(					lowercase    ,		unittest.TestCase    ):





  """simple docstring"""




  _lowerCamelCase						=						(
      (
          FlaxAlbertModel,
          FlaxAlbertForPreTraining,
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxAlbertModelTester(self  )







  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   for model_class_name in self.all_model_classes:
    lowerCamelCase_  				=						model_class_name.from_pretrained("albert-base-v2"  )
    lowerCamelCase_  				=						model(np.ones((1, 1)  )  )
    self.assertIsNotNone(UpperCamelCase  )





@require_flax
class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						FlaxAlbertModel.from_pretrained("albert-base-v2"  )
   lowerCamelCase_  				=						np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
   lowerCamelCase_  				=						np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase  )[0]
   lowerCamelCase_  				=						(1, 11, 768)
   self.assertEqual(output.shape    ,						UpperCamelCase  )
   lowerCamelCase_  				=						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'''

from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput





@dataclass
class   snake_case					(					lowercase    ):





  """simple docstring"""




  _lowerCamelCase						=						None
  _lowerCamelCase						=						None
  _lowerCamelCase						=						None
  _lowerCamelCase						=						None





class   snake_case					(					lowercase    ):





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase=1    ,						UpperCamelCase=0    ,						UpperCamelCase=2    ,						UpperCamelCase=512    ,						UpperCamelCase="cls"    ,						UpperCamelCase=False    ,						UpperCamelCase=True    ,						**UpperCamelCase    ,						):







   """simple docstring"""
   super().__init__(pad_token_id=UpperCamelCase    ,						bos_token_id=UpperCamelCase    ,						eos_token_id=UpperCamelCase    ,						**UpperCamelCase  )
   lowerCamelCase_  				=						project_dim
   lowerCamelCase_  				=						pooler_fn
   lowerCamelCase_  				=						learn_encoder
   lowerCamelCase_  				=						use_attention_mask





class   snake_case					(					lowercase    ):





  """simple docstring"""




  _lowerCamelCase						=						[r"pooler", r"logit_scale"]
  _lowerCamelCase						=						[r"position_ids", r"predictions.decoder.bias"]
  _lowerCamelCase						=						"roberta"
  _lowerCamelCase						=						RobertaSeriesConfig






  def __init__(					self    ,						UpperCamelCase  ):







   """simple docstring"""
   super().__init__(UpperCamelCase  )
   lowerCamelCase_  				=						XLMRobertaModel(UpperCamelCase  )
   lowerCamelCase_  				=						nn.Linear(config.hidden_size    ,						config.project_dim  )
   lowerCamelCase_  				=						getattr(UpperCamelCase    ,						"has_pre_transformation"    ,						UpperCamelCase  )
   if self.has_pre_transformation:
    lowerCamelCase_  				=						nn.Linear(config.hidden_size    ,						config.project_dim  )
    lowerCamelCase_  				=						nn.LayerNorm(config.hidden_size    ,						eps=config.layer_norm_eps  )
   self.post_init()







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







   """simple docstring"""
   lowerCamelCase_  				=						return_dict if return_dict is not None else self.config.use_return_dict

   lowerCamelCase_  				=						self.base_model(
       input_ids=UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						position_ids=UpperCamelCase    ,						head_mask=UpperCamelCase    ,						inputs_embeds=UpperCamelCase    ,						encoder_hidden_states=UpperCamelCase    ,						encoder_attention_mask=UpperCamelCase    ,						output_attentions=UpperCamelCase    ,						output_hidden_states=True if self.has_pre_transformation else output_hidden_states    ,						return_dict=UpperCamelCase    ,						)

   if self.has_pre_transformation:
    lowerCamelCase_  				=						outputs["hidden_states"][-2]
    lowerCamelCase_  				=						self.pre_LN(UpperCamelCase  )
    lowerCamelCase_  				=						self.transformation_pre(UpperCamelCase  )

    return TransformationModelOutput(
        projection_state=UpperCamelCase    ,						last_hidden_state=outputs.last_hidden_state    ,						hidden_states=outputs.hidden_states    ,						attentions=outputs.attentions    ,						)
   else:
    lowerCamelCase_  				=						self.transformation(outputs.last_hidden_state  )
    return TransformationModelOutput(
        projection_state=UpperCamelCase    ,						last_hidden_state=outputs.last_hidden_state    ,						hidden_states=outputs.hidden_states    ,						attentions=outputs.attentions    ,						)



