





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    _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''

 def __init__(self				:   Optional[int]       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   str=13       , UpperCAmelCase_				:   str=7       , UpperCAmelCase_				:   str=True       , UpperCAmelCase_				:   Union[str, Any]=True       , UpperCAmelCase_				:   List[str]=True       , UpperCAmelCase_				:   Tuple=True       , UpperCAmelCase_				:   List[Any]=99       , UpperCAmelCase_				:   Union[str, Any]=24       , UpperCAmelCase_				:   int=2       , UpperCAmelCase_				:   List[Any]=6       , UpperCAmelCase_				:   Tuple=37       , UpperCAmelCase_				:   Tuple="gelu"       , UpperCAmelCase_				:   Dict=0.1       , UpperCAmelCase_				:   Union[str, Any]=0.1       , UpperCAmelCase_				:   Tuple=512       , UpperCAmelCase_				:   Union[str, Any]=16       , UpperCAmelCase_				:   List[Any]=2       , UpperCAmelCase_				:   Tuple=0.02       , UpperCAmelCase_				:   Optional[int]=3       , UpperCAmelCase_				:   Tuple=None       , UpperCAmelCase_				:   int=1_000       , )       ->int:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =parent
   lowerCamelCase__:      Any			  =batch_size
   lowerCamelCase__:      Optional[int]			  =seq_length
   lowerCamelCase__:      Any			  =is_training
   lowerCamelCase__:      Any			  =use_input_mask
   lowerCamelCase__:      Optional[Any]			  =use_token_type_ids
   lowerCamelCase__:      Optional[int]			  =use_labels
   lowerCamelCase__:      str			  =vocab_size
   lowerCamelCase__:      Dict			  =hidden_size
   lowerCamelCase__:      str			  =num_hidden_layers
   lowerCamelCase__:      List[Any]			  =num_attention_heads
   lowerCamelCase__:      Any			  =intermediate_size
   lowerCamelCase__:      Any			  =hidden_act
   lowerCamelCase__:      Optional[int]			  =hidden_dropout_prob
   lowerCamelCase__:      List[Any]			  =attention_probs_dropout_prob
   lowerCamelCase__:      Optional[Any]			  =max_position_embeddings
   lowerCamelCase__:      int			  =type_vocab_size
   lowerCamelCase__:      int			  =type_sequence_label_size
   lowerCamelCase__:      int			  =initializer_range
   lowerCamelCase__:      Optional[int]			  =num_labels
   lowerCamelCase__:      Optional[int]			  =scope
   lowerCamelCase__:      Dict			  =range_bbox

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size)

   lowerCamelCase__:      Optional[Any]			  =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__:      Any			  =bbox[i, j, 3]
         lowerCamelCase__:      Any			  =bbox[i, j, 1]
         lowerCamelCase__:      List[str]			  =t
       if bbox[i, j, 2] < bbox[i, j, 0]:
         lowerCamelCase__:      Dict			  =bbox[i, j, 2]
         lowerCamelCase__:      Any			  =bbox[i, j, 0]
         lowerCamelCase__:      Any			  =t

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

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

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

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

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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->List[Any]:



   '''simple docstring'''



   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 							SCREAMING_SNAKE_CASE_					(self				:   Any       , UpperCAmelCase_				:   Optional[int]       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   Any       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   Union[str, Any]       , )       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =LiltModel(config=UpperCAmelCase_)
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__:      List[Any]			  =model(UpperCAmelCase_       , bbox=UpperCAmelCase_       , attention_mask=UpperCAmelCase_       , token_type_ids=UpperCAmelCase_)
   lowerCamelCase__:      int			  =model(UpperCAmelCase_       , bbox=UpperCAmelCase_       , token_type_ids=UpperCAmelCase_)
   lowerCamelCase__:      Optional[Any]			  =model(UpperCAmelCase_       , bbox=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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Tuple       , )       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      str			  =self.num_labels
   lowerCamelCase__:      Optional[int]			  =LiltForTokenClassification(config=UpperCAmelCase_)
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__:      Tuple			  =model(
       UpperCAmelCase_       , bbox=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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   Optional[int]       , UpperCAmelCase_				:   Union[str, Any]       , UpperCAmelCase_				:   Union[str, Any]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   Any       , )       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =LiltForQuestionAnswering(config=UpperCAmelCase_)
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__:      Tuple			  =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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =self.prepare_config_and_inputs()
   (
       (
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					
   ):      Optional[int]			  =config_and_inputs
   lowerCamelCase__:      str			  ={
       "input_ids": input_ids,
       "bbox": bbox,
       "token_type_ids": token_type_ids,
       "attention_mask": input_mask,
   }
   return config, inputs_dict



@require_torch
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			(
     (
         LiltModel,
         LiltForSequenceClassification,
         LiltForTokenClassification,
         LiltForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 lowercase_		=			(
     {
         "feature-extraction": LiltModel,
         "question-answering": LiltForQuestionAnswering,
         "text-classification": LiltForSequenceClassification,
         "token-classification": LiltForTokenClassification,
         "zero-shot": LiltForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 lowercase_		=			False
 lowercase_		=			False

 def 							SCREAMING_SNAKE_CASE_					(self				:   str       , UpperCAmelCase_				:   List[Any]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   Optional[int]       , UpperCAmelCase_				:   int)       ->Optional[Any]:



   '''simple docstring'''



   return True

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =LiltModelTester(self)
   lowerCamelCase__:      int			  =ConfigTester(self       , config_class=UpperCAmelCase_       , hidden_size=37)

 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->List[str]:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =self.model_tester.prepare_config_and_inputs()
   for type in ["absolute", "relative_key", "relative_key_query"]:
     lowerCamelCase__:      int			  =type
     self.model_tester.create_and_check_model(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Any:



   '''simple docstring'''



   lowerCamelCase__:      int			  =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*UpperCAmelCase_)







 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Tuple:



   '''simple docstring'''



   for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__:      int			  =LiltModel.from_pretrained(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)



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







 '''simple docstring'''

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =LiltModel.from_pretrained("SCUT-DLVCLab/lilt-roberta-en-base").to(UpperCAmelCase_)

   lowerCamelCase__:      Optional[Any]			  =torch.tensor([[1, 2]]       , device=UpperCAmelCase_)
   lowerCamelCase__:      Union[str, Any]			  =torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]       , device=UpperCAmelCase_)

   # forward pass
   with torch.no_grad():
     lowerCamelCase__:      Optional[Any]			  =model(input_ids=UpperCAmelCase_       , bbox=UpperCAmelCase_)

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

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



import os
from tempfile import TemporaryDirectory
from unittest import TestCase

import pytest
from absl.testing import parameterized

from datasets import config
from datasets.arrow_reader import HF_GCP_BASE_URL
from datasets.builder import DatasetBuilder
from datasets.dataset_dict import IterableDatasetDict
from datasets.iterable_dataset import IterableDataset
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.file_utils import cached_path


__lowerCAmelCase			:     Any           =   [
    {'dataset': 'wikipedia', 'config_name': '20220301.de'},
    {'dataset': 'wikipedia', 'config_name': '20220301.en'},
    {'dataset': 'wikipedia', 'config_name': '20220301.fr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.frr'},
    {'dataset': 'wikipedia', 'config_name': '20220301.it'},
    {'dataset': 'wikipedia', 'config_name': '20220301.simple'},
    {'dataset': 'snli', 'config_name': 'plain_text'},
    {'dataset': 'eli5', 'config_name': 'LFQA_reddit'},
    {'dataset': 'wiki40b', 'config_name': 'en'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.compressed'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.nq.no_index'},
    {'dataset': 'wiki_dpr', 'config_name': 'psgs_w100.multiset.no_index'},
    {'dataset': 'natural_questions', 'config_name': 'default'},
]

def a__    (       A_=True					):






  '''simple docstring'''





  if with_config:
    return [
        {
            "testcase_name": d["dataset"] + "/" + d["config_name"],
            "dataset": d["dataset"],
            "config_name": d["config_name"],
        }
        for d in DATASETS_ON_HF_GCP
    ]
  else:
    return [
        {"testcase_name": dataset, "dataset": dataset} for dataset in {d["dataset"] for d in DATASETS_ON_HF_GCP}
    ]




@parameterized.named_parameters(list_datasets_on_hf_gcp_parameters(with_config=_A     )     )
class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  a__           =		None
  a__           =		None
  def 							_lowercase							( self			: List[Any]			,							UpperCamelCase__			: int			,							UpperCamelCase__			: Union[str, Any]    )			->	Tuple:







    """simple docstring"""






    with TemporaryDirectory() as tmp_dir:
      __magic_name__				     =       dataset_module_factory(UpperCamelCase__			,							cache_dir=UpperCamelCase__    )

      __magic_name__				     =       import_main_class(dataset_module.module_path			,							dataset=UpperCamelCase__    )

      __magic_name__				     =       builder_cls(
          cache_dir=UpperCamelCase__			,							config_name=UpperCamelCase__			,							hash=dataset_module.hash			,							)

      __magic_name__				     =       """/""".join(
          [
              HF_GCP_BASE_URL,
              builder_instance._relative_data_dir(with_hash=UpperCamelCase__    ).replace(os.sep			,							"""/"""    ),
              config.DATASET_INFO_FILENAME,
          ]    )
      __magic_name__				     =       cached_path(UpperCamelCase__			,							cache_dir=UpperCamelCase__    )
      self.assertTrue(os.path.exists(UpperCamelCase__    )    )







@pytest.mark.integration
def a__    (       A_					):






  '''simple docstring'''





  __magic_name__				     =       tmp_path_factory.mktemp("""test_hf_gcp"""					) / """test_wikipedia_simple"""
  __magic_name__				     =       dataset_module_factory("""wikipedia""",     cache_dir=A_					)
  __magic_name__				     =       import_main_class(dataset_module.module_path					)
  __magic_name__				     =       builder_cls(
      cache_dir=A_,     config_name="""20220301.frr""",     hash=dataset_module.hash,     )
  # use the HF cloud storage, not the original download_and_prepare that uses apache-beam
  __magic_name__				     =       None
  builder_instance.download_and_prepare()
  __magic_name__				     =       builder_instance.as_dataset()
  assert ds





@pytest.mark.integration
def a__    (       A_					):






  '''simple docstring'''





  __magic_name__				     =       dataset_module_factory("""wikipedia""",     cache_dir=A_					)
  __magic_name__				     =       import_main_class(dataset_module.module_path,     dataset=A_					)
  __magic_name__				     =       builder_cls(
      cache_dir=A_,     config_name="""20220301.frr""",     hash=dataset_module.hash,     )
  __magic_name__				     =       builder_instance.as_streaming_dataset()
  assert ds
  assert isinstance(A_,     A_					)
  assert "train" in ds
  assert isinstance(ds["""train"""],     A_					)
  assert next(iter(ds["""train"""]					)					)

