





import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer

class       A_			(A__     ):
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							tempfile.mkdtemp()
				UpperCAmelCase  						=							5

				# Realm tok
				UpperCAmelCase  						=							[
				    '''[UNK]''',
				    '''[CLS]''',
				    '''[SEP]''',
				    '''[PAD]''',
				    '''[MASK]''',
				    '''test''',
				    '''question''',
				    '''this''',
				    '''is''',
				    '''the''',
				    '''first''',
				    '''second''',
				    '''third''',
				    '''fourth''',
				    '''fifth''',
				    '''record''',
				    '''want''',
				    '''##want''',
				    '''##ed''',
				    '''wa''',
				    '''un''',
				    '''runn''',
				    '''##ing''',
				    ''',''',
				    '''low''',
				    '''lowest''',
				]
				UpperCAmelCase  						=							os.path.join(self.tmpdirname				,   '''realm_tokenizer'''						)
				os.makedirs(_UpperCamelCase				,   exist_ok=_UpperCamelCase						)
				UpperCAmelCase  						=							os.path.join(_UpperCamelCase				,   VOCAB_FILES_NAMES['''vocab_file''']						)
				with open(self.vocab_file				,   '''w'''				,   encoding='''utf-8'''						) as vocab_writer:
							vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]						)						)

				UpperCAmelCase  						=							os.path.join(self.tmpdirname				,   '''realm_block_records'''						)
				os.makedirs(_UpperCamelCase				,   exist_ok=_UpperCamelCase						)
	def 	_lowercase							(       self						):






				'''simple docstring'''







				return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname				,   '''realm_tokenizer'''						)						)
	def 	_lowercase							(       self						):






				'''simple docstring'''







				shutil.rmtree(self.tmpdirname						)
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							RealmConfig(num_block_records=self.num_block_records						)
				return config
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							Dataset.from_dict(
				    {
				        '''id''': ['''0''', '''1'''],
				        '''question''': ['''foo''', '''bar'''],
				        '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
				    }						)
				return dataset
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							np.array(
				    [
				        b'''This is the first record''',
				        b'''This is the second record''',
				        b'''This is the third record''',
				        b'''This is the fourth record''',
				        b'''This is the fifth record''',
				        b'''This is a longer longer longer record''',
				    ]				,   dtype=_UpperCamelCase				,   )
				return block_records
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							RealmRetriever(
				    block_records=self.get_dummy_block_records()				,   tokenizer=self.get_tokenizer()				,   )
				return retriever
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							self.get_config()
				UpperCAmelCase  						=							self.get_dummy_retriever()
				UpperCAmelCase  						=							retriever.tokenizer

				UpperCAmelCase  						=							np.array([0, 3]				,   dtype='''long'''						)
				UpperCAmelCase  						=							tokenizer(['''Test question''']						).input_ids
				UpperCAmelCase  						=							tokenizer(
				    ['''the fourth''']				,   add_special_tokens=_UpperCamelCase				,   return_token_type_ids=_UpperCamelCase				,   return_attention_mask=_UpperCamelCase				,   ).input_ids
				UpperCAmelCase  						=							config.reader_seq_len

				UpperCAmelCase						,     UpperCAmelCase						,     UpperCAmelCase						,     UpperCAmelCase  						=							retriever(
				    _UpperCamelCase				,   _UpperCamelCase				,   answer_ids=_UpperCamelCase				,   max_length=_UpperCamelCase				,   return_tensors='''np'''						)

				self.assertEqual(len(_UpperCamelCase						)				,   2						)
				self.assertEqual(len(_UpperCamelCase						)				,   2						)
				self.assertEqual(len(_UpperCamelCase						)				,   2						)
				self.assertEqual(concat_inputs.input_ids.shape				,   (2, 1_0)						)
				self.assertEqual(concat_inputs.attention_mask.shape				,   (2, 1_0)						)
				self.assertEqual(concat_inputs.token_type_ids.shape				,   (2, 1_0)						)
				self.assertEqual(concat_inputs.special_tokens_mask.shape				,   (2, 1_0)						)
				self.assertEqual(
				    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]						)				,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''first''', '''record''', '''[SEP]''']				,   )
				self.assertEqual(
				    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]						)				,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''fourth''', '''record''', '''[SEP]''']				,   )
	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							self.get_config()
				UpperCAmelCase  						=							self.get_dummy_retriever()
				UpperCAmelCase  						=							retriever.tokenizer

				UpperCAmelCase  						=							np.array([0, 3, 5]				,   dtype='''long'''						)
				UpperCAmelCase  						=							tokenizer(['''Test question''']						).input_ids
				UpperCAmelCase  						=							tokenizer(
				    ['''the fourth''', '''longer longer''']				,   add_special_tokens=_UpperCamelCase				,   return_token_type_ids=_UpperCamelCase				,   return_attention_mask=_UpperCamelCase				,   ).input_ids
				UpperCAmelCase  						=							config.reader_seq_len

				UpperCAmelCase						,     UpperCAmelCase						,     UpperCAmelCase						,     UpperCAmelCase  						=							retriever(
				    _UpperCamelCase				,   _UpperCamelCase				,   answer_ids=_UpperCamelCase				,   max_length=_UpperCamelCase				,   return_tensors='''np'''						)

				self.assertEqual([False, True, True]				,   _UpperCamelCase						)
				self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]				,   _UpperCamelCase						)
				self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]				,   _UpperCamelCase						)







	def 	_lowercase							(       self						):






				'''simple docstring'''







				UpperCAmelCase  						=							self.get_dummy_retriever()
				retriever.save_pretrained(os.path.join(self.tmpdirname				,   '''realm_block_records'''						)						)

				# Test local path
				UpperCAmelCase  						=							retriever.from_pretrained(os.path.join(self.tmpdirname				,   '''realm_block_records'''						)						)
				self.assertEqual(retriever.block_records[0]				,   b'''This is the first record'''						)

				# Test mocked remote path
				with patch('''transformers.models.realm.retrieval_realm.hf_hub_download'''						) as mock_hf_hub_download:
							UpperCAmelCase  						=							os.path.join(
							    os.path.join(self.tmpdirname				,   '''realm_block_records'''						)				,   _REALM_BLOCK_RECORDS_FILENAME						)
							UpperCAmelCase  						=							RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''						)

				self.assertEqual(retriever.block_records[0]				,   b'''This is the first record'''						)










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				,							_UpperCamelCase				,							_UpperCamelCase=13				,							_UpperCamelCase=7				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=True				,							_UpperCamelCase=99				,							_UpperCamelCase=24				,							_UpperCamelCase=2				,							_UpperCamelCase=6				,							_UpperCamelCase=37				,							_UpperCamelCase="gelu"				,							_UpperCamelCase=0.1				,							_UpperCamelCase=0.1				,							_UpperCamelCase=512				,							_UpperCamelCase=16				,							_UpperCamelCase=2				,							_UpperCamelCase=0.02				,							_UpperCamelCase=3				,							_UpperCamelCase=None				,							_UpperCamelCase=1000				,							):
        _UpperCAmelCase           =   parent
        _UpperCAmelCase           =   batch_size
        _UpperCAmelCase           =   seq_length
        _UpperCAmelCase           =   is_training
        _UpperCAmelCase           =   use_input_mask
        _UpperCAmelCase           =   use_token_type_ids
        _UpperCAmelCase           =   use_labels
        _UpperCAmelCase           =   vocab_size
        _UpperCAmelCase           =   hidden_size
        _UpperCAmelCase           =   num_hidden_layers
        _UpperCAmelCase           =   num_attention_heads
        _UpperCAmelCase           =   intermediate_size
        _UpperCAmelCase           =   hidden_act
        _UpperCAmelCase           =   hidden_dropout_prob
        _UpperCAmelCase           =   attention_probs_dropout_prob
        _UpperCAmelCase           =   max_position_embeddings
        _UpperCAmelCase           =   type_vocab_size
        _UpperCAmelCase           =   type_sequence_label_size
        _UpperCAmelCase           =   initializer_range
        _UpperCAmelCase           =   num_labels
        _UpperCAmelCase           =   scope
        _UpperCAmelCase           =   range_bbox





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							self.vocab_size					)

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

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

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

        _UpperCAmelCase           =   None
        _UpperCAmelCase           =   None
        if self.use_labels:
           _UpperCAmelCase           =   ids_tensor([self.batch_size]				,							self.type_sequence_label_size					)
           _UpperCAmelCase           =   ids_tensor([self.batch_size, self.seq_length]				,							self.num_labels					)

        _UpperCAmelCase           =   self.get_config()

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





     def 		UpperCamelCase(     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 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   LiltModel(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   model(_UpperCamelCase				,							bbox=_UpperCamelCase				,							attention_mask=_UpperCamelCase				,							token_type_ids=_UpperCamelCase					)
        _UpperCAmelCase           =   model(_UpperCamelCase				,							bbox=_UpperCamelCase				,							token_type_ids=_UpperCamelCase					)
        _UpperCAmelCase           =   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 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   self.num_labels
        _UpperCAmelCase           =   LiltForTokenClassification(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   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 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							):
        _UpperCAmelCase           =   LiltForQuestionAnswering(config=_UpperCamelCase					)
        model.to(_UpperCamelCase					)
        model.eval()
        _UpperCAmelCase           =   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 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.prepare_config_and_inputs()
        (
            (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      (
            _UpperCAmelCase
        )       ,      
        )           =   config_and_inputs
        _UpperCAmelCase           =   {
            '''input_ids''': input_ids,
            '''bbox''': bbox,
            '''token_type_ids''': token_type_ids,
            '''attention_mask''': input_mask,
        }
        return config, inputs_dict




@require_torch
class 			__UpperCamelCase    (							A__	, A__	, A__	, unittest.TestCase    ):
     __A	:  Dict							  =   (
         (
             LiltModel,
             LiltForSequenceClassification,
             LiltForTokenClassification,
             LiltForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     __A	:  Optional[Any]							  =   (
         {
             """feature-extraction""": LiltModel,
             """question-answering""": LiltForQuestionAnswering,
             """text-classification""": LiltForSequenceClassification,
             """token-classification""": LiltForTokenClassification,
             """zero-shot""": LiltForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __A	:  List[Any]							  =   False
     __A	:  Optional[int]							  =   False





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase				,							_UpperCamelCase					):
        return True





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   LiltModelTester(self					)
        _UpperCAmelCase           =   ConfigTester(self				,							config_class=_UpperCamelCase				,							hidden_size=37					)





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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_UpperCamelCase					)





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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase					)





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*_UpperCamelCase					)





     @slow
     def 		UpperCamelCase(     self					):
        for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           _UpperCAmelCase           =   LiltModel.from_pretrained(_UpperCamelCase					)
           self.assertIsNotNone(_UpperCamelCase					)




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





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''					).to(_UpperCamelCase					)

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

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

        _UpperCAmelCase           =   torch.Size([1, 2, 768]					)
        _UpperCAmelCase           =   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					)					)