


'''simple docstring'''

import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         NezhaForMaskedLM,
         NezhaForMultipleChoice,
         NezhaForNextSentencePrediction,
         NezhaForPreTraining,
         NezhaForQuestionAnswering,
         NezhaForSequenceClassification,
         NezhaForTokenClassification,
         NezhaModel,
     )
     from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST





class   snake_case					:





  """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=128    ,						UpperCamelCase=32    ,						UpperCamelCase=16    ,						UpperCamelCase=2    ,						UpperCamelCase=0.02    ,						UpperCamelCase=3    ,						UpperCamelCase=4    ,						UpperCamelCase=None    ,						):







   """simple docstring"""
   lowerCamelCase_  				=						parent
   lowerCamelCase_  				=						batch_size
   lowerCamelCase_  				=						seq_length
   lowerCamelCase_  				=						is_training
   lowerCamelCase_  				=						use_input_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_labels
   lowerCamelCase_  				=						num_choices
   lowerCamelCase_  				=						scope






  def 						snake_case	(					self  ):







   """simple docstring"""
   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
   if self.use_token_type_ids:
    lowerCamelCase_  				=						ids_tensor([self.batch_size, self.seq_length]    ,						self.type_vocab_size  )

   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_  				=						self.get_config()

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






  def 						snake_case	(					self  ):







   """simple docstring"""
   return NezhaConfig(
       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    ,						)






  def 						snake_case	(					self  ):







   """simple docstring"""
   (
       (
       lowerCamelCase_
   )				,(
       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,
       token_type_ids,
       input_mask,
       sequence_labels,
       token_labels,
       choice_labels,
       encoder_hidden_states,
       encoder_attention_mask,
   )






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







   """simple docstring"""
   lowerCamelCase_  				=						NezhaModel(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase  )
   lowerCamelCase_  				=						model(UpperCamelCase    ,						token_type_ids=UpperCamelCase  )
   lowerCamelCase_  				=						model(UpperCamelCase  )
   self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)  )
   self.parent.assertEqual(result.pooler_output.shape    ,						(self.batch_size, self.hidden_size)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						True
   lowerCamelCase_  				=						NezhaModel(UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						encoder_hidden_states=UpperCamelCase    ,						encoder_attention_mask=UpperCamelCase    ,						)
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						encoder_hidden_states=UpperCamelCase    ,						)
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase  )
   self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)  )
   self.parent.assertEqual(result.pooler_output.shape    ,						(self.batch_size, self.hidden_size)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						NezhaForMaskedLM(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase  )
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						NezhaForNextSentencePrediction(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase    ,						)
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, 2)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						NezhaForPreTraining(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase    ,						next_sentence_label=UpperCamelCase    ,						)
   self.parent.assertEqual(result.prediction_logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)  )
   self.parent.assertEqual(result.seq_relationship_logits.shape    ,						(self.batch_size, 2)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						NezhaForQuestionAnswering(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						start_positions=UpperCamelCase    ,						end_positions=UpperCamelCase    ,						)
   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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.num_labels
   lowerCamelCase_  				=						NezhaForSequenceClassification(UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase  )
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						self.num_labels
   lowerCamelCase_  				=						NezhaForTokenClassification(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase  )
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)  )






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







   """simple docstring"""
   lowerCamelCase_  				=						self.num_choices
   lowerCamelCase_  				=						NezhaForMultipleChoice(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						input_ids.unsqueeze(1  ).expand(-1    ,						self.num_choices    ,						-1  ).contiguous()
   lowerCamelCase_  				=						token_type_ids.unsqueeze(1  ).expand(-1    ,						self.num_choices    ,						-1  ).contiguous()
   lowerCamelCase_  				=						input_mask.unsqueeze(1  ).expand(-1    ,						self.num_choices    ,						-1  ).contiguous()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase    ,						)
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)  )







  def 						snake_case	(					self  ):







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





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





  """simple docstring"""




  _lowerCamelCase						=						(
      (
          NezhaModel,
          NezhaForMaskedLM,
          NezhaForMultipleChoice,
          NezhaForNextSentencePrediction,
          NezhaForPreTraining,
          NezhaForQuestionAnswering,
          NezhaForSequenceClassification,
          NezhaForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase						=						(
      {
          "feature-extraction": NezhaModel,
          "fill-mask": NezhaForMaskedLM,
          "question-answering": NezhaForQuestionAnswering,
          "text-classification": NezhaForSequenceClassification,
          "token-classification": NezhaForTokenClassification,
          "zero-shot": NezhaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  _lowerCamelCase						=						True






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







   """simple docstring"""
   lowerCamelCase_  				=						super()._prepare_for_class(UpperCamelCase    ,						UpperCamelCase    ,						return_labels=UpperCamelCase  )

   if return_labels:
    if model_class in get_values(UpperCamelCase  ):
     lowerCamelCase_  				=						torch.zeros(
         (self.model_tester.batch_size, self.model_tester.seq_length)    ,						dtype=torch.long    ,						device=UpperCamelCase  )
     lowerCamelCase_  				=						torch.zeros(
         self.model_tester.batch_size    ,						dtype=torch.long    ,						device=UpperCamelCase  )
   return inputs_dict






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						NezhaModelTester(self  )
   lowerCamelCase_  				=						ConfigTester(self    ,						config_class=UpperCamelCase    ,						hidden_size=37  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   self.config_tester.run_common_tests()






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   # This regression test was failing with PyTorch < 1.3
   (
       (
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,(
       lowerCamelCase_
   )				,
   )  				=						self.model_tester.prepare_config_and_inputs_for_decoder()

   lowerCamelCase_  				=						None

   self.model_tester.create_and_check_model_as_decoder(
       UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						)






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_next_sequence_prediction(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*UpperCamelCase  )






  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    lowerCamelCase_  				=						NezhaModel.from_pretrained(UpperCamelCase  )
    self.assertIsNotNone(UpperCamelCase  )







  @slow
  @require_torch_gpu
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_				,lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs_for_common()
   for model_class in self.all_model_classes:
    # NezhaForMultipleChoice behaves incorrectly in JIT environments.
    if model_class == NezhaForMultipleChoice:
     return

    lowerCamelCase_  				=						True
    lowerCamelCase_  				=						model_class(config=UpperCamelCase  )

    lowerCamelCase_  				=						self._prepare_for_class(UpperCamelCase    ,						UpperCamelCase  )
    lowerCamelCase_  				=						torch.jit.trace(
        UpperCamelCase    ,						(inputs_dict["input_ids"].to("cpu"  ), inputs_dict["attention_mask"].to("cpu"  ))  )

    with tempfile.TemporaryDirectory() as tmp:
     torch.jit.save(UpperCamelCase    ,						os.path.join(UpperCamelCase    ,						"bert.pt"  )  )
     lowerCamelCase_  				=						torch.jit.load(os.path.join(UpperCamelCase    ,						"bert.pt"  )    ,						map_location=UpperCamelCase  )
     loaded(inputs_dict["input_ids"].to(UpperCamelCase  )    ,						inputs_dict["attention_mask"].to(UpperCamelCase  )  )





@require_torch
class   snake_case					(					unittest.TestCase    ):





  """simple docstring"""




  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						NezhaModel.from_pretrained("sijunhe/nezha-cn-base"  )
   lowerCamelCase_  				=						torch.tensor([[0, 1, 2, 3, 4, 5]]  )
   lowerCamelCase_  				=						torch.tensor([[0, 1, 1, 1, 1, 1]]  )
   with torch.no_grad():
    lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase  )[0]
   lowerCamelCase_  				=						torch.Size((1, 6, 768)  )
   self.assertEqual(output.shape    ,						UpperCamelCase  )
   lowerCamelCase_  				=						torch.tensor([[[0.0_685, 0.2_441, 0.1_102], [0.0_600, 0.1_906, 0.1_349], [0.0_221, 0.0_819, 0.0_586]]]  )

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







  @slow
  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						NezhaForMaskedLM.from_pretrained("sijunhe/nezha-cn-base"  )
   lowerCamelCase_  				=						torch.tensor([[0, 1, 2, 3, 4, 5]]  )
   lowerCamelCase_  				=						torch.tensor([[1, 1, 1, 1, 1, 1]]  )
   with torch.no_grad():
    lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase  )[0]
   lowerCamelCase_  				=						torch.Size((1, 6, 2_1128)  )
   self.assertEqual(output.shape    ,						UpperCamelCase  )
   lowerCamelCase_  				=						torch.tensor(
       [[-2.7_939, -1.7_902, -2.2_189], [-2.8_585, -1.8_908, -2.3_723], [-2.6_499, -1.7_750, -2.2_558]]  )

   self.assertTrue(torch.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    ,						)



