



import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
        LiltForQuestionAnswering,
        LiltForSequenceClassification,
        LiltForTokenClassification,
        LiltModel,
    )
    from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST





class 		UpperCAmelCase						:





   def __init__(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=13       ,     _lowerCAmelCase=7       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=99       ,     _lowerCAmelCase=24       ,     _lowerCAmelCase=2       ,     _lowerCAmelCase=6       ,     _lowerCAmelCase=37       ,     _lowerCAmelCase="gelu"       ,     _lowerCAmelCase=0.1       ,     _lowerCAmelCase=0.1       ,     _lowerCAmelCase=512       ,     _lowerCAmelCase=16       ,     _lowerCAmelCase=2       ,     _lowerCAmelCase=0.02       ,     _lowerCAmelCase=3       ,     _lowerCAmelCase=None       ,     _lowerCAmelCase=1_000       ,     ):
         _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 			=  scope
         _lowerCAmelCase 			=  range_bbox





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size			)

         _lowerCAmelCase 			=  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]:
                           _lowerCAmelCase 			=  bbox[i, j, 3]
                           _lowerCAmelCase 			=  bbox[i, j, 1]
                           _lowerCAmelCase 			=  t
                     if bbox[i, j, 2] < bbox[i, j, 0]:
                           _lowerCAmelCase 			=  bbox[i, j, 2]
                           _lowerCAmelCase 			=  bbox[i, j, 0]
                           _lowerCAmelCase 			=  t

         _lowerCAmelCase 			=  None
         if self.use_input_mask:
               _lowerCAmelCase 			=  ids_tensor([self.batch_size, self.seq_length]       ,     vocab_size=2			)

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

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





   def 	__lowerCAmelCase							(	self			):
         return LiltConfig(
             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 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     ):
         _lowerCAmelCase 			=  LiltModel(config=_lowerCAmelCase			)
         model.to(_lowerCAmelCase			)
         model.eval()
         _lowerCAmelCase 			=  model(_lowerCAmelCase       ,     bbox=_lowerCAmelCase       ,     attention_mask=_lowerCAmelCase       ,     token_type_ids=_lowerCAmelCase			)
         _lowerCAmelCase 			=  model(_lowerCAmelCase       ,     bbox=_lowerCAmelCase       ,     token_type_ids=_lowerCAmelCase			)
         _lowerCAmelCase 			=  model(_lowerCAmelCase       ,     bbox=_lowerCAmelCase			)

         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 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     ):
         _lowerCAmelCase 			=  self.num_labels
         _lowerCAmelCase 			=  LiltForTokenClassification(config=_lowerCAmelCase			)
         model.to(_lowerCAmelCase			)
         model.eval()
         _lowerCAmelCase 			=  model(
             _lowerCAmelCase       ,     bbox=_lowerCAmelCase       ,     attention_mask=_lowerCAmelCase       ,     token_type_ids=_lowerCAmelCase       ,     labels=_lowerCAmelCase			)
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels)			)





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     ):
         _lowerCAmelCase 			=  LiltForQuestionAnswering(config=_lowerCAmelCase			)
         model.to(_lowerCAmelCase			)
         model.eval()
         _lowerCAmelCase 			=  model(
             _lowerCAmelCase       ,     bbox=_lowerCAmelCase       ,     attention_mask=_lowerCAmelCase       ,     token_type_ids=_lowerCAmelCase       ,     start_positions=_lowerCAmelCase       ,     end_positions=_lowerCAmelCase       ,     )
         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 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.prepare_config_and_inputs()
         (
             (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   (
             _lowerCAmelCase
         )			,   
         ) 			=  config_and_inputs
         _lowerCAmelCase 			=  {
             '''input_ids''': input_ids,
             '''bbox''': bbox,
             '''token_type_ids''': token_type_ids,
             '''attention_mask''': input_mask,
         }
         return config, inputs_dict






@require_torch
class 		UpperCAmelCase						(						snake_case_       ,snake_case_       ,snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     (
       (
           LiltModel,
           LiltForSequenceClassification,
           LiltForTokenClassification,
           LiltForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   SCREAMING_SNAKE_CASE__							=     (
       {
           '''feature-extraction''': LiltModel,
           '''question-answering''': LiltForQuestionAnswering,
           '''text-classification''': LiltForSequenceClassification,
           '''token-classification''': LiltForTokenClassification,
           '''zero-shot''': LiltForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   SCREAMING_SNAKE_CASE__							=     False
   SCREAMING_SNAKE_CASE__							=     False





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase       ,     _lowerCAmelCase			):
         return True





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  LiltModelTester(self			)
         _lowerCAmelCase 			=  ConfigTester(self       ,     config_class=_lowerCAmelCase       ,     hidden_size=37			)





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





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
               _lowerCAmelCase 			=  type
               self.model_tester.create_and_check_model(*_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase			)





   @slow
   def 	__lowerCAmelCase							(	self			):
         for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowerCAmelCase 			=  LiltModel.from_pretrained(_lowerCAmelCase			)
               self.assertIsNotNone(_lowerCAmelCase			)






@require_torch
@slow
class 		UpperCAmelCase						(						unittest.TestCase				):





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''			).to(_lowerCAmelCase			)

         _lowerCAmelCase 			=  torch.tensor([[1, 2]]       ,     device=_lowerCAmelCase			)
         _lowerCAmelCase 			=  torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]       ,     device=_lowerCAmelCase			)

         # forward pass
         with torch.no_grad():
               _lowerCAmelCase 			=  model(input_ids=_lowerCAmelCase       ,     bbox=_lowerCAmelCase			)

         _lowerCAmelCase 			=  torch.Size([1, 2, 768]			)
         _lowerCAmelCase 			=  torch.tensor(
             [[-0.0_653, 0.0_950, -0.0_061], [-0.0_545, 0.0_926, -0.0_324]]       ,     device=_lowerCAmelCase       ,     )

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




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch





@require_torch
@require_tf
class 		UpperCAmelCase						(						snake_case_				):





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  SMALL_MODEL_IDENTIFIER
         _lowerCAmelCase 			=  '''pt'''
         _lowerCAmelCase 			=  '''tf'''





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  AutoModel.from_pretrained(self.test_model			)
         model_pt.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self       ,     _lowerCAmelCase			):
         _lowerCAmelCase 			=  TFAutoModel.from_pretrained(self.test_model       ,     from_pt=_lowerCAmelCase			)
         model_tf.save_pretrained(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  '''mock_framework'''

         # Framework provided - return whatever the user provides
         _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model       ,     _lowerCAmelCase			)
         self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # Local checkpoint and framework provided - return provided framework
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     _lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         # PyTorch checkpoint
         with TemporaryDirectory() as local_pt_ckpt:
               self._setup_pt_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # TensorFlow checkpoint
         with TemporaryDirectory() as local_tf_ckpt:
               self._setup_tf_ckpt(_lowerCAmelCase			)
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Invalid local checkpoint
         with TemporaryDirectory() as local_invalid_ckpt:
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(_lowerCAmelCase			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # PyTorch not in environment -> use TensorFlow
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_tf			)

         # Both in environment -> use PyTorch
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)
               self.assertEqual(_lowerCAmelCase       ,     self.framework_pt			)

         # Both not in environment -> raise error
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         _lowerCAmelCase 			=  MagicMock(return_value=_lowerCAmelCase			)
         with patch('''transformers.onnx.features.is_tf_available'''       ,     _lowerCAmelCase			), patch(
             '''transformers.onnx.features.is_torch_available'''       ,     _lowerCAmelCase			):
               with self.assertRaises(_lowerCAmelCase			):
                     _lowerCAmelCase 			=  FeaturesManager.determine_framework(self.test_model			)