






from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
  from transformers import (
      TFAutoModel,
      TFAutoModelForCausalLM,
      TFAutoModelForMaskedLM,
      TFAutoModelForPreTraining,
      TFAutoModelForQuestionAnswering,
      TFAutoModelForSeqaSeqLM,
      TFAutoModelForSequenceClassification,
      TFAutoModelForTableQuestionAnswering,
      TFAutoModelForTokenClassification,
      TFAutoModelWithLMHead,
      TFBertForMaskedLM,
      TFBertForPreTraining,
      TFBertForQuestionAnswering,
      TFBertForSequenceClassification,
      TFBertModel,
      TFFunnelBaseModel,
      TFFunnelModel,
      TFGPTaLMHeadModel,
      TFRobertaForMaskedLM,
      TFTaForConditionalGeneration,
      TFTapasForQuestionAnswering,
  )
  from transformers.models.auto.modeling_tf_auto import (
      TF_MODEL_FOR_CAUSAL_LM_MAPPING,
      TF_MODEL_FOR_MASKED_LM_MAPPING,
      TF_MODEL_FOR_PRETRAINING_MAPPING,
      TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
      TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
      TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
      TF_MODEL_MAPPING,
  )
  from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
  from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST


class    __lowercase   (						lowercase_	):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE     =					"new-model"


if is_tf_available():


  class    __lowercase   (						lowercase_	):







       '''simple docstring'''
       SCREAMING_SNAKE_CASE     =					NewModelConfig


@require_tf
class    __lowercase   (						unittest.TestCase	):







     '''simple docstring'''




     @slow
     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         __A       =						"""bert-base-cased"""
         __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
         self.assertIsNotNone(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

         __A       =						TFAutoModel.from_pretrained(UpperCamelCase_					)
         self.assertIsNotNone(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         __A       =						"""bert-base-cased"""
         __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
         self.assertIsNotNone(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

         __A       =						TFAutoModelForPreTraining.from_pretrained(UpperCamelCase_					)
         self.assertIsNotNone(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      List[Any]					):


         """simple docstring"""



         for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForCausalLM.from_pretrained(UpperCamelCase_					)
             __A     ,						__A       =						TFAutoModelForCausalLM.from_pretrained(UpperCamelCase_			,						output_loading_info=UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      Union[str, Any]					):


         """simple docstring"""



         for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelWithLMHead.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      int					):


         """simple docstring"""



         for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForMaskedLM.from_pretrained(UpperCamelCase_					)
             __A     ,						__A       =						TFAutoModelForMaskedLM.from_pretrained(UpperCamelCase_			,						output_loading_info=UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForSeqaSeqLM.from_pretrained(UpperCamelCase_					)
             __A     ,						__A       =						TFAutoModelForSeqaSeqLM.from_pretrained(UpperCamelCase_			,						output_loading_info=UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      Tuple					):


         """simple docstring"""



         for model_name in ["bert-base-uncased"]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForSequenceClassification.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      int					):


         """simple docstring"""



         for model_name in ["bert-base-uncased"]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForQuestionAnswering.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     @require_tensorflow_probability
     def 		lowerCAmelCase_  (					self							:      Tuple					):


         """simple docstring"""



         for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
             __A       =						AutoConfig.from_pretrained(UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

             __A       =						TFAutoModelForTableQuestionAnswering.from_pretrained(UpperCamelCase_					)
             __A     ,						__A       =						TFAutoModelForTableQuestionAnswering.from_pretrained(
                 UpperCamelCase_			,						output_loading_info=UpperCamelCase_					)
             self.assertIsNotNone(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     def 		lowerCAmelCase_  (					self							:      Any					):


         """simple docstring"""



         __A       =						TFAutoModelWithLMHead.from_pretrained(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)
         self.assertEqual(model.num_parameters()			,						14_410					)
         self.assertEqual(model.num_parameters(only_trainable=UpperCamelCase_					)			,						14_410					)




     def 		lowerCAmelCase_  (					self							:      Optional[Any]					):


         """simple docstring"""



         __A       =						TFAutoModelWithLMHead.from_pretrained(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)
         self.assertEqual(model.num_parameters()			,						14_410					)
         self.assertEqual(model.num_parameters(only_trainable=UpperCamelCase_					)			,						14_410					)




     def 		lowerCAmelCase_  (					self							:      Optional[Any]					):


         """simple docstring"""



         __A       =						TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

         __A       =						copy.deepcopy(model.config					)
         __A       =						["""FunnelBaseModel"""]
         __A       =						TFAutoModel.from_config(UpperCamelCase_					)
         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

         with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(UpperCamelCase_					)
             __A       =						TFAutoModel.from_pretrained(UpperCamelCase_					)
             self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         try:
             AutoConfig.register("""new-model"""			,						UpperCamelCase_					)

             __A       =						[
                 TFAutoModel,
                 TFAutoModelForCausalLM,
                 TFAutoModelForMaskedLM,
                 TFAutoModelForPreTraining,
                 TFAutoModelForQuestionAnswering,
                 TFAutoModelForSequenceClassification,
                 TFAutoModelForTokenClassification,
             ]

             for auto_class in auto_classes:
                 with self.subTest(auto_class.__name__					):
                     # Wrong config class will raise an error
                     with self.assertRaises(UpperCamelCase_					):
                         auto_class.register(UpperCamelCase_			,						UpperCamelCase_					)
                     auto_class.register(UpperCamelCase_			,						UpperCamelCase_					)
                     # Trying to register something existing in the Transformers library will raise an error
                     with self.assertRaises(UpperCamelCase_					):
                         auto_class.register(UpperCamelCase_			,						UpperCamelCase_					)

                     # Now that the config is registered, it can be used as any other config with the auto-API
                     __A       =						BertModelTester(self					).get_config()
                     __A       =						NewModelConfig(**tiny_config.to_dict()					)
                     __A       =						auto_class.from_config(UpperCamelCase_					)
                     self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

                     with tempfile.TemporaryDirectory() as tmp_dir:
                         model.save_pretrained(UpperCamelCase_					)
                         __A       =						auto_class.from_pretrained(UpperCamelCase_					)
                         self.assertIsInstance(UpperCamelCase_			,						UpperCamelCase_					)

         finally:
             if "new-model" in CONFIG_MAPPING._extra_content:
                 del CONFIG_MAPPING._extra_content["new-model"]
             for mapping in (
                 TF_MODEL_MAPPING,
                 TF_MODEL_FOR_PRETRAINING_MAPPING,
                 TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                 TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                 TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                 TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                 TF_MODEL_FOR_MASKED_LM_MAPPING,
             ):
                 if NewModelConfig in mapping._extra_content:
                     del mapping._extra_content[NewModelConfig]




     def 		lowerCAmelCase_  (					self							:      Union[str, Any]					):


         """simple docstring"""



         with self.assertRaisesRegex(
             UpperCamelCase_			,						"""bert-base is not a local folder and is not a valid model identifier"""					):
             __A       =						TFAutoModel.from_pretrained("""bert-base"""					)




     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         with self.assertRaisesRegex(
             UpperCamelCase_			,						R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""					):
             __A       =						TFAutoModel.from_pretrained(UpperCamelCase_			,						revision="""aaaaaa"""					)




     def 		lowerCAmelCase_  (					self							:      List[str]					):


         """simple docstring"""



         with self.assertRaisesRegex(
             UpperCamelCase_			,						"""hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""			,						):
             __A       =						TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""					)




     def 		lowerCAmelCase_  (					self							:      Any					):


         """simple docstring"""



         with self.assertRaisesRegex(UpperCamelCase_			,						"""Use `from_pt=True` to load this model"""					):
             __A       =						TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""					)





     def 		lowerCAmelCase_  (					self							:      Optional[int]					):


         """simple docstring"""



         __A       =						TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""					)
         with RequestCounter() as counter:
             __A       =						TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""					)
             self.assertEqual(counter.get_request_count			,						0					)
             self.assertEqual(counter.head_request_count			,						1					)
             self.assertEqual(counter.other_request_count			,						0					)

         # With a sharded checkpoint
         __A       =						TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""					)
         with RequestCounter() as counter:
             __A       =						TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""					)
             self.assertEqual(counter.get_request_count			,						0					)
             self.assertEqual(counter.head_request_count			,						1					)
             self.assertEqual(counter.other_request_count			,						0					)












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    __lowercase   (						lowercase_	):







     '''simple docstring'''




     def 		lowerCAmelCase_  (					self							:      int					):


         """simple docstring"""



         __A       =						tempfile.mkdtemp()
         __A       =						5

         # Realm tok
         __A       =						[
             """[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""",
         ]
         __A       =						os.path.join(self.tmpdirname			,						"""realm_tokenizer"""					)
         os.makedirs(UpperCamelCase_			,						exist_ok=UpperCamelCase_					)
         __A       =						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]					)					)

         __A       =						os.path.join(self.tmpdirname			,						"""realm_block_records"""					)
         os.makedirs(UpperCamelCase_			,						exist_ok=UpperCamelCase_					)




     def 		lowerCAmelCase_  (					self							:      List[Any]					):


         """simple docstring"""



         return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,						"""realm_tokenizer"""					)					)




     def 		lowerCAmelCase_  (					self							:      Union[str, Any]					):


         """simple docstring"""



         shutil.rmtree(self.tmpdirname					)




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						RealmConfig(num_block_records=self.num_block_records					)
         return config




     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



         __A       =						Dataset.from_dict(
             {
                 """id""": ["""0""", """1"""],
                 """question""": ["""foo""", """bar"""],
                 """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
             }					)
         return dataset




     def 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						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 		lowerCAmelCase_  (					self							:      List[Any]					):


         """simple docstring"""



         __A       =						RealmRetriever(
             block_records=self.get_dummy_block_records()			,						tokenizer=self.get_tokenizer()			,						)
         return retriever




     def 		lowerCAmelCase_  (					self							:      Dict					):


         """simple docstring"""



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

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

         __A     ,						__A     ,						__A     ,						__A       =						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, 10)					)
         self.assertEqual(concat_inputs.attention_mask.shape			,						(2, 10)					)
         self.assertEqual(concat_inputs.token_type_ids.shape			,						(2, 10)					)
         self.assertEqual(concat_inputs.special_tokens_mask.shape			,						(2, 10)					)
         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 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



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

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

         __A     ,						__A     ,						__A     ,						__A       =						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 		lowerCAmelCase_  (					self							:      str					):


         """simple docstring"""



         __A       =						self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname			,						"""realm_block_records"""					)					)

         # Test local path
         __A       =						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:
             __A       =						os.path.join(
                 os.path.join(self.tmpdirname			,						"""realm_block_records"""					)			,						_REALM_BLOCK_RECORDS_FILENAME					)
             __A       =						RealmRetriever.from_pretrained("""google/realm-cc-news-pretrained-openqa"""					)

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




