






"""simple docstring"""






import argparse
import json
import os
from collections import OrderedDict

import torch

from transformers import LukeConfig, LukeForMaskedLM, MLukeTokenizer, XLMRobertaTokenizer
from transformers.tokenization_utils_base import AddedToken



@torch.no_grad()
def 					_A			(__a    ,				__a    ,				__a    ,				__a    ,				__a							) ->       int:
    """simple docstring"""






    with open(__a							) as metadata_file:
        SCREAMING_SNAKE_CASE_      :  int    						=		json.load(__a							)
    SCREAMING_SNAKE_CASE_      :  str    						=		LukeConfig(use_entity_aware_attention=__a    ,				**metadata['''model_config''']							)

    # Load in the weights from the checkpoint_path
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		torch.load(__a    ,				map_location='''cpu'''							)['''module''']

    # Load the entity vocab file
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		load_original_entity_vocab(__a							)
    # add an entry for [MASK2]
    SCREAMING_SNAKE_CASE_      :  Any    						=		max(entity_vocab.values()							) + 1
    config.entity_vocab_size += 1

    SCREAMING_SNAKE_CASE_      :  List[str]    						=		XLMRobertaTokenizer.from_pretrained(metadata['''model_config''']['''bert_model_name''']							)

    # Add special tokens to the token vocabulary for downstream tasks
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		AddedToken('''<ent>'''    ,				lstrip=__a    ,				rstrip=__a							)
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		AddedToken('''<ent2>'''    ,				lstrip=__a    ,				rstrip=__a							)
    tokenizer.add_special_tokens({'''additional_special_tokens''': [entity_token_a, entity_token_a]}							)
    config.vocab_size += 2

    print(f'Saving tokenizer to {pytorch_dump_folder_path}'							)
    tokenizer.save_pretrained(__a							)
    with open(os.path.join(__a    ,				'''tokenizer_config.json'''							)    ,				'''r'''							) as f:
        SCREAMING_SNAKE_CASE_      :  str    						=		json.load(__a							)
    SCREAMING_SNAKE_CASE_      :  str    						=		'''MLukeTokenizer'''
    with open(os.path.join(__a    ,				'''tokenizer_config.json'''							)    ,				'''w'''							) as f:
        json.dump(__a    ,				__a							)

    with open(os.path.join(__a    ,				MLukeTokenizer.vocab_files_names['''entity_vocab_file''']							)    ,				'''w'''							) as f:
        json.dump(__a    ,				__a							)

    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		MLukeTokenizer.from_pretrained(__a							)

    # Initialize the embeddings of the special tokens
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		tokenizer.convert_tokens_to_ids(['''@''']							)[0]
    SCREAMING_SNAKE_CASE_      :  str    						=		tokenizer.convert_tokens_to_ids(['''#''']							)[0]

    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		state_dict['''embeddings.word_embeddings.weight''']
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		word_emb[ent_init_index].unsqueeze(0							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		word_emb[enta_init_index].unsqueeze(0							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		torch.cat([word_emb, ent_emb, enta_emb]							)
    # add special tokens for 'entity_predictions.bias'
    for bias_name in ["lm_head.decoder.bias", "lm_head.bias"]:
        SCREAMING_SNAKE_CASE_      :  Tuple    						=		state_dict[bias_name]
        SCREAMING_SNAKE_CASE_      :  Dict    						=		decoder_bias[ent_init_index].unsqueeze(0							)
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		decoder_bias[enta_init_index].unsqueeze(0							)
        SCREAMING_SNAKE_CASE_      :  List[Any]    						=		torch.cat([decoder_bias, ent_decoder_bias, enta_decoder_bias]							)

    # Initialize the query layers of the entity-aware self-attention mechanism
    for layer_index in range(config.num_hidden_layers							):
        for matrix_name in ["query.weight", "query.bias"]:
            SCREAMING_SNAKE_CASE_      :  List[Any]    						=		f'encoder.layer.{layer_index}.attention.self.'
            SCREAMING_SNAKE_CASE_      :  Tuple    						=		state_dict[prefix + matrix_name]
            SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		state_dict[prefix + matrix_name]
            SCREAMING_SNAKE_CASE_      :  Dict    						=		state_dict[prefix + matrix_name]

    # Initialize the embedding of the [MASK2] entity using that of the [MASK] entity for downstream tasks
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		state_dict['''entity_embeddings.entity_embeddings.weight''']
    SCREAMING_SNAKE_CASE_      :  Any    						=		entity_emb[entity_vocab['''[MASK]''']].unsqueeze(0							)
    SCREAMING_SNAKE_CASE_      :  str    						=		torch.cat([entity_emb, entity_mask_emb]							)
    # add [MASK2] for 'entity_predictions.bias'
    SCREAMING_SNAKE_CASE_      :  str    						=		state_dict['''entity_predictions.bias''']
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		entity_prediction_bias[entity_vocab['''[MASK]''']].unsqueeze(0							)
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		torch.cat([entity_prediction_bias, entity_mask_bias]							)

    SCREAMING_SNAKE_CASE_      :  Any    						=		LukeForMaskedLM(config=__a							).eval()

    state_dict.pop('''entity_predictions.decoder.weight'''							)
    state_dict.pop('''lm_head.decoder.weight'''							)
    state_dict.pop('''lm_head.decoder.bias'''							)
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		OrderedDict()
    for key, value in state_dict.items():
        if not (key.startswith('''lm_head'''							) or key.startswith('''entity_predictions'''							)):
            SCREAMING_SNAKE_CASE_      :  List[Any]    						=		state_dict[key]
        else:
            SCREAMING_SNAKE_CASE_      :  Tuple    						=		state_dict[key]

    SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  List[Any]    						=		model.load_state_dict(__a    ,				strict=__a							)

    if set(__a							) != {"luke.embeddings.position_ids"}:
        raise ValueError(f'Unexpected unexpected_keys: {unexpected_keys}'							)
    if set(__a							) != {
        "lm_head.decoder.weight",
        "lm_head.decoder.bias",
        "entity_predictions.decoder.weight",
    }:
        raise ValueError(f'Unexpected missing_keys: {missing_keys}'							)

    model.tie_weights()
    assert (model.luke.embeddings.word_embeddings.weight == model.lm_head.decoder.weight).all()
    assert (model.luke.entity_embeddings.entity_embeddings.weight == model.entity_predictions.decoder.weight).all()

    # Check outputs
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		MLukeTokenizer.from_pretrained(__a    ,				task='''entity_classification'''							)

    SCREAMING_SNAKE_CASE_      :  Tuple    						=		'''ISO 639-3 uses the code fas for the dialects spoken across Iran and アフガニスタン (Afghanistan).'''
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		(0, 9)
    SCREAMING_SNAKE_CASE_      :  Any    						=		tokenizer(__a    ,				entity_spans=[span]    ,				return_tensors='''pt'''							)

    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		model(**__a							)

    # Verify word hidden states
    if model_size == "large":
        raise NotImplementedError
    else:  # base
        SCREAMING_SNAKE_CASE_      :  str    						=		torch.Size((1, 33, 7_68)							)
        SCREAMING_SNAKE_CASE_      :  List[str]    						=		torch.tensor([[0.08_92, 0.05_96, -0.28_19], [0.01_34, 0.11_99, 0.05_73], [-0.01_69, 0.09_27, 0.06_44]]							)

    if not (outputs.last_hidden_state.shape == expected_shape):
        raise ValueError(
            f'Outputs.last_hidden_state.shape is {outputs.last_hidden_state.shape}, Expected shape is {expected_shape}'							)
    if not torch.allclose(outputs.last_hidden_state[0, :3, :3]    ,				__a    ,				atol=1e-4							):
        raise ValueError

    # Verify entity hidden states
    if model_size == "large":
        raise NotImplementedError
    else:  # base
        SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		torch.Size((1, 1, 7_68)							)
        SCREAMING_SNAKE_CASE_      :  Dict    						=		torch.tensor([[-0.14_82, 0.06_09, 0.03_22]]							)

    if not (outputs.entity_last_hidden_state.shape == expected_shape):
        raise ValueError(
            f'Outputs.entity_last_hidden_state.shape is {outputs.entity_last_hidden_state.shape}, Expected shape is'
            f' {expected_shape}'							)
    if not torch.allclose(outputs.entity_last_hidden_state[0, :3, :3]    ,				__a    ,				atol=1e-4							):
        raise ValueError

    # Verify masked word/entity prediction
    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		MLukeTokenizer.from_pretrained(__a							)
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		'''Tokyo is the capital of <mask>.'''
    SCREAMING_SNAKE_CASE_      :  Dict    						=		(24, 30)
    SCREAMING_SNAKE_CASE_      :  List[str]    						=		tokenizer(__a    ,				entity_spans=[span]    ,				return_tensors='''pt'''							)

    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		model(**__a							)

    SCREAMING_SNAKE_CASE_      :  Tuple    						=		encoding['''input_ids'''][0].tolist()
    SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		input_ids.index(tokenizer.convert_tokens_to_ids('''<mask>'''							)							)
    SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		outputs.logits[0][mask_position_id].argmax(dim=-1							)
    assert "Japan" == tokenizer.decode(__a							)

    SCREAMING_SNAKE_CASE_      :  List[str]    						=		outputs.entity_logits[0][0].argmax().item()
    SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
        entity for entity, entity_id in tokenizer.entity_vocab.items() if entity_id == predicted_entity_id
    ]
    assert [e for e in multilingual_predicted_entities if e.startswith('''en:'''							)][0] == "en:Japan"

    # Finally, save our PyTorch model and tokenizer
    print('''Saving PyTorch model to {}'''.format(__a							)							)
    model.save_pretrained(__a							)



def 					_A			(__a							) ->       int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  List[Any]    						=		['''[MASK]''', '''[PAD]''', '''[UNK]''']

    SCREAMING_SNAKE_CASE_      :  List[str]    						=		[json.loads(__a							) for line in open(__a							)]

    SCREAMING_SNAKE_CASE_      :  int    						=		{}
    for entry in data:
        SCREAMING_SNAKE_CASE_      :  Tuple    						=		entry['''id''']
        for entity_name, language in entry["entities"]:
            if entity_name in SPECIAL_TOKENS:
                SCREAMING_SNAKE_CASE_      :  str    						=		entity_id
                break
            SCREAMING_SNAKE_CASE_      :  List[str]    						=		f'{language}:{entity_name}'
            SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		entity_id
    return new_mapping


if __name__ == "__main__":
  UpperCAmelCase_   : List[str]   			= argparse.ArgumentParser()
  # Required parameters
  parser.add_argument("""--checkpoint_path""", type=str, help="""Path to a pytorch_model.bin file.""")
  parser.add_argument(
      """--metadata_path""", default=None, type=str, help="""Path to a metadata.json file, defining the configuration."""
  )
  parser.add_argument(
      """--entity_vocab_path""",
      default=None,
      type=str,
      help="""Path to an entity_vocab.tsv file, containing the entity vocabulary.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to where to dump the output PyTorch model."""
  )
  parser.add_argument(
      """--model_size""", default="""base""", type=str, choices=["""base""", """large"""], help="""Size of the model to be converted."""
  )
  UpperCAmelCase_   : List[Any]   			= parser.parse_args()
  convert_luke_checkpoint(
      args.checkpoint_path,
      args.metadata_path,
      args.entity_vocab_path,
      args.pytorch_dump_folder_path,
      args.model_size,
  )










"""simple docstring"""






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






      '''simple docstring'''

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




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		parent
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		batch_size
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		seq_length
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		is_training
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		use_input_mask
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		use_token_type_ids
          SCREAMING_SNAKE_CASE_      :  int    						=		use_labels
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		vocab_size
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		hidden_size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		num_hidden_layers
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		num_attention_heads
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		intermediate_size
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		hidden_act
          SCREAMING_SNAKE_CASE_      :  int    						=		hidden_dropout_prob
          SCREAMING_SNAKE_CASE_      :  str    						=		attention_probs_dropout_prob
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		max_position_embeddings
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		type_vocab_size
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		type_sequence_label_size
          SCREAMING_SNAKE_CASE_      :  Any    						=		initializer_range
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		num_labels
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		scope
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		range_bbox
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size)

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

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

          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		None
          if self.use_token_type_ids:
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		ids_tensor([self.batch_size, self.seq_length]						,	self.type_vocab_size)

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

          SCREAMING_SNAKE_CASE_      :  Any    						=		self.get_config()

          return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''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     :  Optional[int]						,	lowercase_     :  Optional[Any]						,	lowercase_     :  Optional[Any]						,	lowercase_     :  str						,	lowercase_     :  Optional[Any]						,	lowercase_     :  int						,	lowercase_     :  Optional[Any]						,	lowercase_     :  int						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		LiltModel(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		model(lowercase_						,	bbox=lowercase_						,	attention_mask=lowercase_						,	token_type_ids=lowercase_)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		model(lowercase_						,	bbox=lowercase_						,	token_type_ids=lowercase_)
          SCREAMING_SNAKE_CASE_      :  int    						=		model(lowercase_						,	bbox=lowercase_)

          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     :  int						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  Optional[Any]						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  List[Any]						,	lowercase_     :  int						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.num_labels
          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		LiltForTokenClassification(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		model(
              lowercase_						,	bbox=lowercase_						,	attention_mask=lowercase_						,	token_type_ids=lowercase_						,	labels=lowercase_)
          self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels))
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  str						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  Any						,	lowercase_     :  Dict						,	lowercase_     :  List[str]						,	lowercase_     :  List[Any]						,	lowercase_     :  str						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		LiltForQuestionAnswering(config=lowercase_)
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		model(
              lowercase_						,	bbox=lowercase_						,	attention_mask=lowercase_						,	token_type_ids=lowercase_						,	start_positions=lowercase_						,	end_positions=lowercase_						,	)
          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     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  str    						=		self.prepare_config_and_inputs()
          (
              (
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			(
              SCREAMING_SNAKE_CASE_
          )						,			
          )      :  List[str]    						=		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 ):






      '''simple docstring'''

      __UpperCamelCase	 =  (
          (
              LiltModel,
              LiltForSequenceClassification,
              LiltForTokenClassification,
              LiltForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      __UpperCamelCase	 =  (
          {
              "feature-extraction": LiltModel,
              "question-answering": LiltForQuestionAnswering,
              "text-classification": LiltForSequenceClassification,
              "token-classification": LiltForTokenClassification,
              "zero-shot": LiltForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      __UpperCamelCase	 =  False
      __UpperCamelCase	 =  False
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  List[str]						,	lowercase_     :  Optional[int]						,	lowercase_     :  Optional[int]						,	lowercase_     :  str						,	lowercase_     :  str):




          '''simple docstring'''


          return True
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		LiltModelTester(self)
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		ConfigTester(self						,	config_class=lowercase_						,	hidden_size=37)
      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          self.config_tester.run_common_tests()
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
              SCREAMING_SNAKE_CASE_      :  Dict    						=		type
              self.model_tester.create_and_check_model(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Tuple):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*lowercase_)



      @slow
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		LiltModel.from_pretrained(lowercase_)
              self.assertIsNotNone(lowercase_)




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






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base''').to(lowercase_)

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

          # forward pass
          with torch.no_grad():
              SCREAMING_SNAKE_CASE_      :  Dict    						=		model(input_ids=lowercase_						,	bbox=lowercase_)

          SCREAMING_SNAKE_CASE_      :  str    						=		torch.Size([1, 2, 768])
          SCREAMING_SNAKE_CASE_      :  Dict    						=		torch.tensor(
              [[-0.06_53, 0.09_50, -0.00_61], [-0.05_45, 0.09_26, -0.03_24]]						,	device=lowercase_						,	)

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


