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 		_lowerCamelCase		:
    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=1000 , )       ->					str:
     SCREAMING_SNAKE_CASE__:   str=      parent
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      batch_size
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      seq_length
     SCREAMING_SNAKE_CASE__:   Any=      is_training
     SCREAMING_SNAKE_CASE__:   Dict=      use_input_mask
     SCREAMING_SNAKE_CASE__:   str=      use_token_type_ids
     SCREAMING_SNAKE_CASE__:   Any=      use_labels
     SCREAMING_SNAKE_CASE__:   Dict=      vocab_size
     SCREAMING_SNAKE_CASE__:   Tuple=      hidden_size
     SCREAMING_SNAKE_CASE__:   Optional[int]=      num_hidden_layers
     SCREAMING_SNAKE_CASE__:   List[str]=      num_attention_heads
     SCREAMING_SNAKE_CASE__:   Optional[int]=      intermediate_size
     SCREAMING_SNAKE_CASE__:   str=      hidden_act
     SCREAMING_SNAKE_CASE__:   List[Any]=      hidden_dropout_prob
     SCREAMING_SNAKE_CASE__:   Tuple=      attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE__:   str=      max_position_embeddings
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      type_vocab_size
     SCREAMING_SNAKE_CASE__:   str=      type_sequence_label_size
     SCREAMING_SNAKE_CASE__:   Optional[int]=      initializer_range
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      num_labels
     SCREAMING_SNAKE_CASE__:   Union[str, Any]=      scope
     SCREAMING_SNAKE_CASE__:   Any=      range_bbox
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Any=      ids_tensor([self.batch_size, self.seq_length] , self.vocab_size   )

     SCREAMING_SNAKE_CASE__:   Tuple=      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]:
        SCREAMING_SNAKE_CASE__:   str=      bbox[i, j, 3]
        SCREAMING_SNAKE_CASE__:   str=      bbox[i, j, 1]
        SCREAMING_SNAKE_CASE__:   Union[str, Any]=      t
       if bbox[i, j, 2] < bbox[i, j, 0]:
        SCREAMING_SNAKE_CASE__:   Optional[Any]=      bbox[i, j, 2]
        SCREAMING_SNAKE_CASE__:   Any=      bbox[i, j, 0]
        SCREAMING_SNAKE_CASE__:   Union[str, Any]=      t

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

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

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

     SCREAMING_SNAKE_CASE__:   List[Any]=      self.get_config()

     return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     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 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , )       ->					Optional[Any]:
     SCREAMING_SNAKE_CASE__:   int=      LiltModel(config=lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()
     SCREAMING_SNAKE_CASE__:   List[Any]=      model(lowerCAmelCase , bbox=lowerCAmelCase , attention_mask=lowerCAmelCase , token_type_ids=lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      model(lowerCAmelCase , bbox=lowerCAmelCase , token_type_ids=lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   Optional[int]=      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 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , )       ->					Optional[Any]:
     SCREAMING_SNAKE_CASE__:   str=      self.num_labels
     SCREAMING_SNAKE_CASE__:   Dict=      LiltForTokenClassification(config=lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()
     SCREAMING_SNAKE_CASE__:   List[Any]=      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 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   int=      LiltForQuestionAnswering(config=lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()
     SCREAMING_SNAKE_CASE__:   Optional[int]=      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 						UpperCamelCase_				(  self   )       ->					List[str]:
     SCREAMING_SNAKE_CASE__:   Dict=      self.prepare_config_and_inputs()
     (
         SCREAMING_SNAKE_CASE__	
     ):   Optional[int]=      config_and_inputs
     SCREAMING_SNAKE_CASE__:   str=      {
         '''input_ids''': input_ids,
         '''bbox''': bbox,
         '''token_type_ids''': token_type_ids,
         '''attention_mask''': input_mask,
     }
     return config, inputs_dict






@require_torch
class 		_lowerCamelCase		(   UpperCamelCase_    ,     UpperCamelCase_    ,     UpperCamelCase_    ,     unittest.TestCase    ):
    __a  =      (
        (
            LiltModel,
            LiltForSequenceClassification,
            LiltForTokenClassification,
            LiltForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    __a  =      (
        {
            "feature-extraction": LiltModel,
            "question-answering": LiltForQuestionAnswering,
            "text-classification": LiltForSequenceClassification,
            "token-classification": LiltForTokenClassification,
            "zero-shot": LiltForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __a  =      False
    __a  =      False
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase   )       ->					List[str]:
     return True
    def 						UpperCamelCase_				(  self   )       ->					int:
     SCREAMING_SNAKE_CASE__:   str=      LiltModelTester(self   )
     SCREAMING_SNAKE_CASE__:   int=      ConfigTester(self , config_class=lowerCAmelCase , hidden_size=37   )
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     self.config_tester.run_common_tests()
    def 						UpperCamelCase_				(  self   )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   Any=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowerCAmelCase   )
    def 						UpperCamelCase_				(  self   )       ->					Any:
     SCREAMING_SNAKE_CASE__:   Any=      self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
      SCREAMING_SNAKE_CASE__:   str=      type
      self.model_tester.create_and_check_model(*lowerCAmelCase   )
    def 						UpperCamelCase_				(  self   )       ->					List[str]:
     SCREAMING_SNAKE_CASE__:   List[Any]=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase   )
    def 						UpperCamelCase_				(  self   )       ->					Any:
     SCREAMING_SNAKE_CASE__:   str=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase   )



    @slow
    def 						UpperCamelCase_				(  self   )       ->					Optional[int]:
     for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      SCREAMING_SNAKE_CASE__:   Union[str, Any]=      LiltModel.from_pretrained(lowerCAmelCase   )
      self.assertIsNotNone(lowerCAmelCase   )






@require_torch
@slow
class 		_lowerCamelCase		(   unittest.TestCase    ):
    def 						UpperCamelCase_				(  self   )       ->					Any:
     SCREAMING_SNAKE_CASE__:   str=      LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''   ).to(lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   int=      torch.tensor([[1, 2]] , device=lowerCAmelCase   )
     SCREAMING_SNAKE_CASE__:   str=      torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]] , device=lowerCAmelCase   )

     # forward pass
     with torch.no_grad():
      SCREAMING_SNAKE_CASE__:   Tuple=      model(input_ids=lowerCAmelCase , bbox=lowerCAmelCase   )

     SCREAMING_SNAKE_CASE__:   Any=      torch.Size([1, 2, 768]   )
     SCREAMING_SNAKE_CASE__:   str=      torch.tensor(
         [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]] , device=lowerCAmelCase , )

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


import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

lowercase_     :   Optional[Any] 				=				[
    'cross_validation.py',
    'gradient_accumulation.py',
    'local_sgd.py',
    'multi_process_metrics.py',
    'memory.py',
    'automatic_gradient_accumulation.py',
    'fsdp_with_peak_mem_tracking.py',
    'deepspeed_with_config_support.py',
    'megatron_lm_gpt_pretraining.py',
]






class 		_lowerCamelCase		(   unittest.TestCase    ):
    def 						UpperCamelCase_				(  self , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase = None , lowerCAmelCase = None   )       ->					int:
     SCREAMING_SNAKE_CASE__:   Optional[int]=      None
     SCREAMING_SNAKE_CASE__:   Dict=      os.path.abspath(os.path.join('''examples''' , '''by_feature'''   )   )
     SCREAMING_SNAKE_CASE__:   List[Any]=      os.path.abspath('''examples'''   )
     for item in os.listdir(lowerCAmelCase   ):
      if item not in EXCLUDE_EXAMPLES:
       SCREAMING_SNAKE_CASE__:   Tuple=      os.path.join(lowerCAmelCase , lowerCAmelCase   )
       if os.path.isfile(lowerCAmelCase   ) and ".py" in item_path:
        with self.subTest(
            tested_script=lowerCAmelCase , feature_script=lowerCAmelCase , tested_section='''main()''' if parser_only else '''training_function()''' , ):
         SCREAMING_SNAKE_CASE__:   List[Any]=      compare_against_test(
             os.path.join(lowerCAmelCase , lowerCAmelCase   ) , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase   )
         SCREAMING_SNAKE_CASE__:   str=      '''\n'''.join(lowerCAmelCase   )
         if special_strings is not None:
          for string in special_strings:
           SCREAMING_SNAKE_CASE__:   Union[str, Any]=      diff.replace(lowerCAmelCase , ''''''   )
         self.assertEqual(lowerCAmelCase , ''''''   )
    def 						UpperCamelCase_				(  self   )       ->					Optional[Any]:
     self.one_complete_example('''complete_nlp_example.py''' , lowerCAmelCase   )
     self.one_complete_example('''complete_nlp_example.py''' , lowerCAmelCase   )



    def 						UpperCamelCase_				(  self   )       ->					Tuple:
     SCREAMING_SNAKE_CASE__:   List[Any]=      os.path.abspath(os.path.join('''examples''' , '''cv_example.py'''   )   )
     SCREAMING_SNAKE_CASE__:   Dict=      [
         ''' ''' * 16 + '''{\n\n''',
         ''' ''' * 20 + '''"accuracy": eval_metric["accuracy"],\n\n''',
         ''' ''' * 20 + '''"f1": eval_metric["f1"],\n\n''',
         ''' ''' * 20 + '''"train_loss": total_loss.item() / len(train_dataloader),\n\n''',
         ''' ''' * 20 + '''"epoch": epoch,\n\n''',
         ''' ''' * 16 + '''},\n\n''',
         ''' ''' * 16 + '''step=epoch,\n''',
         ''' ''' * 12,
         ''' ''' * 8 + '''for step, batch in enumerate(active_dataloader):\n''',
     ]
     self.one_complete_example('''complete_cv_example.py''' , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase   )
     self.one_complete_example('''complete_cv_example.py''' , lowerCAmelCase , lowerCAmelCase , lowerCAmelCase   )






@mock.patch.dict(os.environ    ,     {"TESTING_MOCKED_DATALOADERS": "1"}    )
class 		_lowerCamelCase		(   UpperCamelCase_    ):
    __a  =      False
    @classmethod
    def 						UpperCamelCase_				(  cls   )       ->					List[Any]:
     super().setUpClass()
     SCREAMING_SNAKE_CASE__:   List[Any]=      tempfile.mkdtemp()
     SCREAMING_SNAKE_CASE__:   Tuple=      os.path.join(cls._tmpdir , '''default_config.yml'''   )

     write_basic_config(save_location=cls.configPath   )
     SCREAMING_SNAKE_CASE__:   Dict=      ['''accelerate''', '''launch''', '''--config_file''', cls.configPath]
    @classmethod
    def 						UpperCamelCase_				(  cls   )       ->					List[Any]:
     super().tearDownClass()
     shutil.rmtree(cls._tmpdir   )
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   Dict=      f'\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        '.split()
     run_command(self._launch_args + testargs   )
     self.assertTrue(os.path.exists(os.path.join(self.tmpdir , '''epoch_0'''   )   )   )
    def 						UpperCamelCase_				(  self   )       ->					Dict:
     SCREAMING_SNAKE_CASE__:   List[Any]=      f'\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        '.split()
     SCREAMING_SNAKE_CASE__:   List[Any]=      run_command(self._launch_args + testargs   )
     self.assertTrue(os.path.exists(os.path.join(self.tmpdir , '''step_2'''   )   )   )
    def 						UpperCamelCase_				(  self   )       ->					Dict:
     SCREAMING_SNAKE_CASE__:   Tuple=      f'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir , "epoch_0"   )}\n        '.split()
     SCREAMING_SNAKE_CASE__:   Dict=      run_command(self._launch_args + testargs , return_stdout=lowerCAmelCase   )
     self.assertNotIn('''epoch 0:''' , lowerCAmelCase   )
     self.assertIn('''epoch 1:''' , lowerCAmelCase   )
    def 						UpperCamelCase_				(  self   )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   Dict=      f'\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir , "step_2"   )}\n        '.split()
     SCREAMING_SNAKE_CASE__:   Dict=      run_command(self._launch_args + testargs , return_stdout=lowerCAmelCase   )
     if torch.cuda.is_available():
      SCREAMING_SNAKE_CASE__:   List[str]=      torch.cuda.device_count()
     else:
      SCREAMING_SNAKE_CASE__:   List[Any]=      1
     if num_processes > 1:
      self.assertNotIn('''epoch 0:''' , lowerCAmelCase   )
      self.assertIn('''epoch 1:''' , lowerCAmelCase   )
     else:
      self.assertIn('''epoch 0:''' , lowerCAmelCase   )
      self.assertIn('''epoch 1:''' , lowerCAmelCase   )
    @slow
    def 						UpperCamelCase_				(  self   )       ->					str:
     SCREAMING_SNAKE_CASE__:   List[str]=      '''
        examples/by_feature/cross_validation.py
        --num_folds 2
        '''.split()
     with mock.patch.dict(os.environ , {'''TESTING_MOCKED_DATALOADERS''': '''0'''}   ):
      SCREAMING_SNAKE_CASE__:   Any=      run_command(self._launch_args + testargs , return_stdout=lowerCAmelCase   )
      SCREAMING_SNAKE_CASE__:   Dict=      re.findall('''({.+})''' , lowerCAmelCase   )
      SCREAMING_SNAKE_CASE__:   int=      [r for r in results if '''accuracy''' in r][-1]
      SCREAMING_SNAKE_CASE__:   List[str]=      ast.literal_eval(lowerCAmelCase   )
      self.assertGreaterEqual(results['''accuracy'''] , 0.75   )
    def 						UpperCamelCase_				(  self   )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   Optional[Any]=      ['''examples/by_feature/multi_process_metrics.py''']
     run_command(self._launch_args + testargs   )
    @require_trackers
    @mock.patch.dict(os.environ , {'''WANDB_MODE''': '''offline'''}   )
    def 						UpperCamelCase_				(  self   )       ->					int:
     with tempfile.TemporaryDirectory() as tmpdir:
      SCREAMING_SNAKE_CASE__:   Optional[int]=      f'\n            examples/by_feature/tracking.py\n            --with_tracking\n            --project_dir {tmpdir}\n            '.split()
      run_command(self._launch_args + testargs   )
      self.assertTrue(os.path.exists(os.path.join(lowerCAmelCase , '''tracking'''   )   )   )
    def 						UpperCamelCase_				(  self   )       ->					int:
     SCREAMING_SNAKE_CASE__:   int=      ['''examples/by_feature/gradient_accumulation.py''']
     run_command(self._launch_args + testargs   )



    def 						UpperCamelCase_				(  self   )       ->					List[Any]:
     SCREAMING_SNAKE_CASE__:   List[str]=      ['''examples/by_feature/local_sgd.py''']
     run_command(self._launch_args + testargs   )

