

import unittest

from transformers import LayoutLMConfig, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          LayoutLMForMaskedLM,
          LayoutLMForQuestionAnswering,
          LayoutLMForSequenceClassification,
          LayoutLMForTokenClassification,
          LayoutLMModel,
      )


class      lowerCamelCase_ :




 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.0_2  ,  _UpperCAmelCase=3  ,  _UpperCAmelCase=4  ,  _UpperCAmelCase=None  ,  _UpperCAmelCase=1_000  ,  ):
   __UpperCAmelCase	:  Union[str, Any]		=     parent
   __UpperCAmelCase	:  List[Any]		=     batch_size
   __UpperCAmelCase	:  Tuple		=     seq_length
   __UpperCAmelCase	:  List[str]		=     is_training
   __UpperCAmelCase	:  int		=     use_input_mask
   __UpperCAmelCase	:  Dict		=     use_token_type_ids
   __UpperCAmelCase	:  Optional[Any]		=     use_labels
   __UpperCAmelCase	:  str		=     vocab_size
   __UpperCAmelCase	:  str		=     hidden_size
   __UpperCAmelCase	:  Optional[int]		=     num_hidden_layers
   __UpperCAmelCase	:  str		=     num_attention_heads
   __UpperCAmelCase	:  int		=     intermediate_size
   __UpperCAmelCase	:  List[str]		=     hidden_act
   __UpperCAmelCase	:  Union[str, Any]		=     hidden_dropout_prob
   __UpperCAmelCase	:  Tuple		=     attention_probs_dropout_prob
   __UpperCAmelCase	:  str		=     max_position_embeddings
   __UpperCAmelCase	:  Tuple		=     type_vocab_size
   __UpperCAmelCase	:  Optional[Any]		=     type_sequence_label_size
   __UpperCAmelCase	:  str		=     initializer_range
   __UpperCAmelCase	:  Union[str, Any]		=     num_labels
   __UpperCAmelCase	:  int		=     num_choices
   __UpperCAmelCase	:  str		=     scope
   __UpperCAmelCase	:  Optional[Any]		=     range_bbox




 def     a       (			self							):
   __UpperCAmelCase	:  List[str]		=     ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size							)

   __UpperCAmelCase	:  List[str]		=     ids_tensor([self.batch_size, self.seq_length, 4]  ,  self.range_bbox							)
   # Ensure that bbox is legal
   for i in range(bbox.shape[0]							):
     for j in range(bbox.shape[1]							):
       if bbox[i, j, 3] < bbox[i, j, 1]:
         __UpperCAmelCase	:  Any		=     bbox[i, j, 3]
         __UpperCAmelCase	:  Optional[int]		=     bbox[i, j, 1]
         __UpperCAmelCase	:  Dict		=     t
       if bbox[i, j, 2] < bbox[i, j, 0]:
         __UpperCAmelCase	:  Optional[Any]		=     bbox[i, j, 2]
         __UpperCAmelCase	:  Tuple		=     bbox[i, j, 0]
         __UpperCAmelCase	:  Dict		=     t

   __UpperCAmelCase	:  Optional[int]		=     None
   if self.use_input_mask:
     __UpperCAmelCase	:  Any		=     ids_tensor([self.batch_size, self.seq_length]  ,  vocab_size=2							)

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

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

   __UpperCAmelCase	:  Optional[Any]		=     self.get_config()

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




 def     a       (			self							):
   return LayoutLMConfig(
       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     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  Tuple		=     LayoutLMModel(config=_UpperCAmelCase							)
   model.to(_UpperCAmelCase							)
   model.eval()
   __UpperCAmelCase	:  Dict		=     model(_UpperCAmelCase  ,  _UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase							)
   __UpperCAmelCase	:  Optional[int]		=     model(_UpperCAmelCase  ,  _UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase							)
   __UpperCAmelCase	:  Union[str, Any]		=     model(_UpperCAmelCase  ,  _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     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  Union[str, Any]		=     LayoutLMForMaskedLM(config=_UpperCAmelCase							)
   model.to(_UpperCAmelCase							)
   model.eval()
   __UpperCAmelCase	:  List[str]		=     model(_UpperCAmelCase  ,  _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     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  str		=     self.num_labels
   __UpperCAmelCase	:  str		=     LayoutLMForSequenceClassification(_UpperCAmelCase							)
   model.to(_UpperCAmelCase							)
   model.eval()
   __UpperCAmelCase	:  str		=     model(
       _UpperCAmelCase  ,  _UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase  ,  labels=_UpperCAmelCase							)
   self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)							)




 def     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  int		=     self.num_labels
   __UpperCAmelCase	:  int		=     LayoutLMForTokenClassification(config=_UpperCAmelCase							)
   model.to(_UpperCAmelCase							)
   model.eval()
   __UpperCAmelCase	:  str		=     model(_UpperCAmelCase  ,  _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     a       (			self  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase  ,  _UpperCAmelCase							):
   __UpperCAmelCase	:  Optional[Any]		=     LayoutLMForQuestionAnswering(config=_UpperCAmelCase							)
   model.to(_UpperCAmelCase							)
   model.eval()
   __UpperCAmelCase	:  Optional[Any]		=     model(
       _UpperCAmelCase  ,  bbox=_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     a       (			self							):
   __UpperCAmelCase	:  int		=     self.prepare_config_and_inputs()
   (
       (
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			(
       __UpperCAmelCase
   ),			
   )	:  Union[str, Any]		=     config_and_inputs
   __UpperCAmelCase	:  Tuple		=     {
       'input_ids': input_ids,
       'bbox': bbox,
       'token_type_ids': token_type_ids,
       'attention_mask': input_mask,
   }
   return config, inputs_dict



@require_torch
class      lowerCamelCase_ (			_lowerCAmelCase	,				_lowerCAmelCase	,				unittest.TestCase):
 _lowerCAmelCase									=	(
     (
         LayoutLMModel,
         LayoutLMForMaskedLM,
         LayoutLMForSequenceClassification,
         LayoutLMForTokenClassification,
         LayoutLMForQuestionAnswering,
     )
     if is_torch_available()
     else None
 )
 _lowerCAmelCase									=	(
     {
         '''document-question-answering''': LayoutLMForQuestionAnswering,
         '''feature-extraction''': LayoutLMModel,
         '''fill-mask''': LayoutLMForMaskedLM,
         '''text-classification''': LayoutLMForSequenceClassification,
         '''token-classification''': LayoutLMForTokenClassification,
         '''zero-shot''': LayoutLMForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 _lowerCAmelCase									=	True




 def     a       (			self							):
   __UpperCAmelCase	:  int		=     LayoutLMModelTester(self							)
   __UpperCAmelCase	:  Tuple		=     ConfigTester(self  ,  config_class=_UpperCAmelCase  ,  hidden_size=37							)




 def     a       (			self							):
   self.config_tester.run_common_tests()




 def     a       (			self							):
   __UpperCAmelCase	:  List[Any]		=     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_UpperCAmelCase							)




 def     a       (			self							):
   __UpperCAmelCase	:  int		=     self.model_tester.prepare_config_and_inputs()
   for type in ["absolute", "relative_key", "relative_key_query"]:
     __UpperCAmelCase	:  int		=     type
     self.model_tester.create_and_check_model(*_UpperCAmelCase							)




 def     a       (			self							):
   __UpperCAmelCase	:  Dict		=     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase							)




 def     a       (			self							):
   __UpperCAmelCase	:  str		=     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*_UpperCAmelCase							)




 def     a       (			self							):
   __UpperCAmelCase	:  List[Any]		=     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase							)




 def     a       (			self							):
   __UpperCAmelCase	:  List[str]		=     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*_UpperCAmelCase							)




 @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'							)
 def     a       (			self							):
   pass





def  _snake_case		()							->		Tuple:
  # Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
  # fmt: off
  __UpperCAmelCase	:  List[str]		=     torch.tensor([[101,1019,1014,1016,1037,1_2849,4747,1004,1_4246,2278,5439,4524,5002,2930,2193,2930,4341,3208,1005,1055,2171,2848,1_1300,3531,102],[101,4070,4034,7020,1024,3058,1015,1013,2861,1013,6070,1_9274,2772,6205,2_7814,1_6147,1_6147,4343,2047,1_0283,1_0969,1_4389,1012,2338,102]]    ,       device=snake_case_  )  # noqa: E231
  __UpperCAmelCase	:  Union[str, Any]		=     torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]    ,       device=snake_case_  )  # noqa: E231
  __UpperCAmelCase	:  List[Any]		=     torch.tensor([[[0,0,0,0],[423,237,440,251],[427,272,441,287],[419,115,437,129],[961,885,992,912],[256,38,330,58],[256,38,330,58],[336,42,353,57],[360,39,401,56],[360,39,401,56],[411,39,471,59],[479,41,528,59],[533,39,630,60],[67,113,134,131],[141,115,209,132],[68,149,133,166],[141,149,187,164],[195,148,287,165],[195,148,287,165],[195,148,287,165],[295,148,349,165],[441,149,492,166],[497,149,546,164],[64,201,125,218],[1000,1000,1000,1000]],[[0,0,0,0],[662,150,754,166],[665,199,742,211],[519,213,554,228],[519,213,554,228],[134,433,187,454],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[130,467,204,480],[314,469,376,482],[504,684,582,706],[941,825,973,900],[941,825,973,900],[941,825,973,900],[941,825,973,900],[610,749,652,765],[130,659,168,672],[176,657,237,672],[238,657,312,672],[443,653,628,672],[443,653,628,672],[716,301,825,317],[1000,1000,1000,1000]]]    ,       device=snake_case_  )  # noqa: E231
  __UpperCAmelCase	:  Optional[Any]		=     torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]    ,       device=snake_case_  )  # noqa: E231
  # these are sequence labels (i.e. at the token level)
  __UpperCAmelCase	:  Union[str, Any]		=     torch.tensor([[-100,10,10,10,9,1,-100,7,7,-100,7,7,4,2,5,2,8,8,-100,-100,5,0,3,2,-100],[-100,12,12,12,-100,12,10,-100,-100,-100,-100,10,12,9,-100,-100,-100,10,10,10,9,12,-100,10,-100]]    ,       device=snake_case_  )  # noqa: E231
  # fmt: on

  return input_ids, attention_mask, bbox, token_type_ids, labels







@require_torch
class      lowerCamelCase_ (			unittest.TestCase):




 @slow
 def     a       (			self							):
   __UpperCAmelCase	:  Tuple		=     LayoutLMModel.from_pretrained('microsoft/layoutlm-base-uncased'							).to(_UpperCAmelCase							)

   __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  Union[str, Any]		=     prepare_layoutlm_batch_inputs()

   # forward pass
   __UpperCAmelCase	:  Any		=     model(input_ids=_UpperCAmelCase  ,  bbox=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase							)

   # test the sequence output on [0, :3, :3]
   __UpperCAmelCase	:  str		=     torch.tensor(
       [[0.1_7_8_5, -0.1_9_4_7, -0.0_4_2_5], [-0.3_2_5_4, -0.2_8_0_7, 0.2_5_5_3], [-0.5_3_9_1, -0.3_3_2_2, 0.3_3_6_4]]  ,  device=_UpperCAmelCase  ,  )

   self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]  ,  _UpperCAmelCase  ,  atol=1E-3							)							)

   # test the pooled output on [1, :3]
   __UpperCAmelCase	:  Any		=     torch.tensor([-0.6_5_8_0, -0.0_2_1_4, 0.8_5_5_2]  ,  device=_UpperCAmelCase							)

   self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]  ,  _UpperCAmelCase  ,  atol=1E-3							)							)




 @slow
 def     a       (			self							):
   # initialize model with randomly initialized sequence classification head
   __UpperCAmelCase	:  List[str]		=     LayoutLMForSequenceClassification.from_pretrained('microsoft/layoutlm-base-uncased'  ,  num_labels=2							).to(
       _UpperCAmelCase							)

   __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  Dict		=     prepare_layoutlm_batch_inputs()

   # forward pass
   __UpperCAmelCase	:  Optional[Any]		=     model(
       input_ids=_UpperCAmelCase  ,  bbox=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase  ,  labels=torch.tensor([1, 1]  ,  device=_UpperCAmelCase							)  ,  )

   # test whether we get a loss as a scalar
   __UpperCAmelCase	:  str		=     outputs.loss
   __UpperCAmelCase	:  List[str]		=     torch.Size([]							)
   self.assertEqual(loss.shape  ,  _UpperCAmelCase							)

   # test the shape of the logits
   __UpperCAmelCase	:  Union[str, Any]		=     outputs.logits
   __UpperCAmelCase	:  int		=     torch.Size((2, 2)							)
   self.assertEqual(logits.shape  ,  _UpperCAmelCase							)




 @slow
 def     a       (			self							):
   # initialize model with randomly initialized token classification head
   __UpperCAmelCase	:  List[Any]		=     LayoutLMForTokenClassification.from_pretrained('microsoft/layoutlm-base-uncased'  ,  num_labels=13							).to(
       _UpperCAmelCase							)

   __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  List[str]		=     prepare_layoutlm_batch_inputs()

   # forward pass
   __UpperCAmelCase	:  Optional[int]		=     model(
       input_ids=_UpperCAmelCase  ,  bbox=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase  ,  labels=_UpperCAmelCase							)

   # test the loss calculation to be around 2.65
   # expected_loss = torch.tensor(2.65, device=torch_device)

   # The loss is currently somewhat random and can vary between 0.1-0.3 atol.
   # self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

   # test the shape of the logits
   __UpperCAmelCase	:  List[str]		=     outputs.logits
   __UpperCAmelCase	:  Optional[Any]		=     torch.Size((2, 25, 13)							)
   self.assertEqual(logits.shape  ,  _UpperCAmelCase							)




 @slow
 def     a       (			self							):
   # initialize model with randomly initialized token classification head
   __UpperCAmelCase	:  Any		=     LayoutLMForQuestionAnswering.from_pretrained('microsoft/layoutlm-base-uncased'							).to(_UpperCAmelCase							)

   __UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase,			__UpperCAmelCase	:  Optional[Any]		=     prepare_layoutlm_batch_inputs()

   # forward pass
   __UpperCAmelCase	:  Optional[Any]		=     model(input_ids=_UpperCAmelCase  ,  bbox=_UpperCAmelCase  ,  attention_mask=_UpperCAmelCase  ,  token_type_ids=_UpperCAmelCase							)

   # test the shape of the logits
   __UpperCAmelCase	:  List[str]		=     torch.Size((2, 25)							)
   self.assertEqual(outputs.start_logits.shape  ,  _UpperCAmelCase							)
   self.assertEqual(outputs.end_logits.shape  ,  _UpperCAmelCase							)









import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  import jax
  import jax.numpy as jnp
  import numpy as np

  from transformers.modeling_flax_pytorch_utils import (
      convert_pytorch_state_dict_to_flax,
      load_flax_weights_in_pytorch_model,
  )
  from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
  import torch

@require_flax
class   lowerCamelCase						:

       def __init__(      self						,    a__						,    a__=14						,    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__=0.02						,    a__=None						,    ):






        '''simple docstring'''






        snake_case :			Any							=   parent
        snake_case :			str							=   batch_size
        snake_case :			str							=   seq_length
        snake_case :			str							=   is_training
        snake_case :			List[str]							=   use_input_mask
        snake_case :			str							=   use_labels
        snake_case :			Union[str, Any]							=   vocab_size
        snake_case :			Dict							=   d_model
        snake_case :			List[str]							=   num_hidden_layers
        snake_case :			Optional[int]							=   num_attention_heads
        snake_case :			Optional[Any]							=   ffn_dim
        snake_case :			Dict							=   activation_function
        snake_case :			List[str]							=   activation_dropout
        snake_case :			Any							=   attention_dropout
        snake_case :			str							=   max_position_embeddings
        snake_case :			Optional[int]							=   initializer_range
        snake_case :			Union[str, Any]							=   None
        snake_case :			int							=   0
        snake_case :			int							=   2
        snake_case :			Union[str, Any]							=   1

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   np.clip(ids_tensor([self.batch_size, self.seq_length]						,    self.vocab_size       )						,    3						,    self.vocab_size       )

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

        snake_case :			Optional[Any]							=   XGLMConfig(
            vocab_size=self.vocab_size						,    d_model=self.hidden_size						,    num_layers=self.num_hidden_layers						,    attention_heads=self.num_attention_heads						,    ffn_dim=self.ffn_dim						,    activation_function=self.activation_function						,    activation_dropout=self.activation_dropout						,    attention_dropout=self.attention_dropout						,    max_position_embeddings=self.max_position_embeddings						,    initializer_range=self.initializer_range						,    use_cache=a__						,    bos_token_id=self.bos_token_id						,    eos_token_id=self.eos_token_id						,    pad_token_id=self.pad_token_id						,    )

        return (config, input_ids, input_mask)

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






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

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case	,       snake_case :			List[str]							=   self.prepare_config_and_inputs()

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

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )

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






        '''simple docstring'''






        snake_case :			Optional[Any]							=   20
        snake_case :			List[Any]							=   model_class_name(a__       )

        snake_case :			List[Any]							=   model.init_cache(input_ids.shape[0]						,    a__       )
        snake_case :			Optional[int]							=   jnp.ones((input_ids.shape[0], max_decoder_length)						,    dtype='''i4'''       )

        snake_case :			Optional[Any]							=   jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)       )
        snake_case :			Tuple							=   model(
            input_ids[:, :-1]						,    attention_mask=a__						,    past_key_values=a__						,    position_ids=a__						,    )

        snake_case :			Dict							=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype='''i4'''       )
        snake_case :			str							=   model(
            input_ids[:, -1:]						,    attention_mask=a__						,    past_key_values=outputs_cache.past_key_values						,    position_ids=a__						,    )

        snake_case :			Dict							=   model(a__       )

        snake_case :			Union[str, Any]							=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
        self.parent.assertTrue(diff < 1e-3						,    msg=f'''Max diff is {diff}'''       )



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






        '''simple docstring'''






        snake_case :			Optional[int]							=   20
        snake_case :			Dict							=   model_class_name(a__       )

        snake_case :			Dict							=   jnp.concatenate(
            [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])       )]						,    axis=-1						,    )

        snake_case :			List[str]							=   model.init_cache(input_ids.shape[0]						,    a__       )
        snake_case :			Optional[int]							=   jnp.broadcast_to(
            jnp.arange(input_ids.shape[-1] - 1       )[None, :]						,    (input_ids.shape[0], input_ids.shape[-1] - 1)       )

        snake_case :			Optional[int]							=   model(
            input_ids[:, :-1]						,    attention_mask=a__						,    past_key_values=a__						,    position_ids=a__						,    )
        snake_case :			int							=   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,    dtype='''i4'''       )
        snake_case :			Optional[int]							=   model(
            input_ids[:, -1:]						,    past_key_values=outputs_cache.past_key_values						,    attention_mask=a__						,    position_ids=a__						,    )

        snake_case :			Optional[Any]							=   model(a__						,    attention_mask=a__       )
        snake_case :			Optional[int]							=   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])       )       )
        self.parent.assertTrue(diff < 1e-3						,    msg=f'''Max diff is {diff}'''       )







@require_sentencepiece
@require_flax
class   lowerCamelCase						(	SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      unittest.TestCase  ):
       _UpperCamelCase     		=   (FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
       _UpperCamelCase     		=   (FlaxXGLMForCausalLM,) if is_flax_available() else ()

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   FlaxXGLMModelTester(self       )

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case	,       snake_case	,       snake_case :			str							=   self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward(a__						,    a__						,    a__						,    a__       )

       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case	,       snake_case	,       snake_case :			int							=   self.model_tester.prepare_config_and_inputs()
         self.model_tester.check_use_cache_forward_with_attn_mask(
             a__						,    a__						,    a__						,    a__       )

       @slow
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case :			Any							=   XGLMTokenizer.from_pretrained('''XGLM'''						,    padding_side='''left'''       )
        snake_case :			Optional[int]							=   tokenizer(['''Hello this is a long string''', '''Hey''']						,    return_tensors='''np'''						,    padding=a__						,    truncation=a__       )

        snake_case :			List[Any]							=   FlaxXGLMForCausalLM.from_pretrained('''facebook/xglm-564M'''       )
        snake_case :			Optional[int]							=   1
        snake_case :			str							=   False

        snake_case :			str							=   jax.jit(model.generate       )

        snake_case :			Dict							=   jit_generate(inputs['''input_ids''']						,    attention_mask=inputs['''attention_mask''']       ).sequences

        snake_case :			Optional[Any]							=   tokenizer.batch_decode(a__						,    skip_special_tokens=a__       )

        snake_case :			Tuple							=   [
            '''Hello this is a long string of questions, but I\'m not sure if I\'m''',
            '''Hey, I\'m a newbie to the forum and I\'''',
        ]

        self.assertListEqual(a__						,    a__       )

       @is_pt_flax_cross_test
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case :			Union[str, Any]							=   self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__       ):
          # prepare inputs
          snake_case :			Union[str, Any]							=   self._prepare_for_class(a__						,    a__       )
          snake_case :			int							=   {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          snake_case :			Any							=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          snake_case :			Optional[Any]							=   getattr(a__						,    a__       )

          snake_case	,       snake_case :			Optional[Any]							=   pt_inputs['''input_ids'''].shape
          snake_case :			Tuple							=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)       )
          for batch_idx, start_index in enumerate(a__       ):
           snake_case :			Optional[Any]							=   0
           snake_case :			Optional[int]							=   1
           snake_case :			Optional[Any]							=   0
           snake_case :			Optional[Any]							=   1
          snake_case :			Dict							=   pt_model_class(a__       ).eval()
          # Flax models don't use the `use_cache` option and cache is not returned as a default.
          # So we disable `use_cache` here for PyTorch model.
          snake_case :			Any							=   False
          snake_case :			str							=   model_class(a__						,    dtype=jnp.floataa       )
          snake_case :			Optional[int]							=   convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,    a__       )
          snake_case :			Optional[int]							=   fx_state

          with torch.no_grad():
           snake_case :			int							=   pt_model(**a__       ).to_tuple()

          snake_case :			Optional[int]							=   fx_model(**a__       ).to_tuple()
          self.assertEqual(len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

          with tempfile.TemporaryDirectory() as tmpdirname:
           pt_model.save_pretrained(a__       )
           snake_case :			Union[str, Any]							=   model_class.from_pretrained(a__						,    from_pt=a__       )

          snake_case :			List[Any]							=   fx_model_loaded(**a__       ).to_tuple()
          self.assertEqual(
              len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output_loaded, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output_loaded[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

       @is_pt_flax_cross_test
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        snake_case	,       snake_case :			List[Any]							=   self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
         with self.subTest(model_class.__name__       ):
          # prepare inputs
          snake_case :			str							=   self._prepare_for_class(a__						,    a__       )
          snake_case :			List[Any]							=   {k: torch.tensor(v.tolist()       ) for k, v in prepared_inputs_dict.items()}

          # load corresponding PyTorch class
          snake_case :			Any							=   model_class.__name__[4:]  # Skip the "Flax" at the beginning
          snake_case :			Union[str, Any]							=   getattr(a__						,    a__       )

          snake_case :			List[str]							=   pt_model_class(a__       ).eval()
          snake_case :			Tuple							=   False
          snake_case :			Optional[int]							=   model_class(a__						,    dtype=jnp.floataa       )

          snake_case :			Any							=   load_flax_weights_in_pytorch_model(a__						,    fx_model.params       )
          snake_case	,       snake_case :			Union[str, Any]							=   pt_inputs['''input_ids'''].shape
          snake_case :			List[Any]							=   np.random.randint(0						,    seq_length - 1						,    size=(batch_size,)       )
          for batch_idx, start_index in enumerate(a__       ):
           snake_case :			Union[str, Any]							=   0
           snake_case :			Tuple							=   1
           snake_case :			int							=   0
           snake_case :			Optional[int]							=   1

          # make sure weights are tied in PyTorch
          pt_model.tie_weights()

          with torch.no_grad():
           snake_case :			Optional[int]							=   pt_model(**a__       ).to_tuple()

          snake_case :			Any							=   fx_model(**a__       ).to_tuple()
          self.assertEqual(len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )

          with tempfile.TemporaryDirectory() as tmpdirname:
           fx_model.save_pretrained(a__       )
           snake_case :			int							=   pt_model_class.from_pretrained(a__						,    from_flax=a__       )

          with torch.no_grad():
           snake_case :			int							=   pt_model_loaded(**a__       ).to_tuple()

          self.assertEqual(
              len(a__       )						,    len(a__       )						,    '''Output lengths differ between Flax and PyTorch'''       )
          for fx_output, pt_output in zip(a__						,    a__       ):
           self.assert_almost_equals(fx_output[:, -1]						,    pt_output[:, -1].numpy()						,    4e-2       )



       @slow
       def 					UpperCAmelCase					(      self       ):






        '''simple docstring'''






        for model_class_name in self.all_model_classes:
         snake_case :			str							=   model_class_name.from_pretrained('''facebook/xglm-564M'''       )
         snake_case :			List[Any]							=   model(np.ones((1, 1)       )       )
         self.assertIsNotNone(a__       )






