




'''simple docstring'''




import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow

class     lowerCAmelCase_ :







  '''simple docstring'''







  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[Any]    ,		_UpperCAmelCase     :						Dict    ,		_UpperCAmelCase     :						Union[str, Any]    ,		_UpperCAmelCase     :						int			):






   """simple docstring"""






   return None

class     lowerCAmelCase_ :







  '''simple docstring'''







  def 				SCREAMING_SNAKE_CASE__      (    self     :						Dict    ,		_UpperCAmelCase     :						Optional[Any]    ,		_UpperCAmelCase     :						Union[str, Any]    ,		_UpperCAmelCase     :						Optional[int]    ,		_UpperCAmelCase     :						Dict			):






   """simple docstring"""






   return None

class     lowerCAmelCase_ (			unittest.TestCase ):







  '''simple docstring'''







  lowerCAmelCase_					:       str					  =    [
      # (model_name, model_kwargs)
      ("""bert-base-cased""", {}),
      ("""gpt2""", {"""use_cache""": False}),  # We don't support exporting GPT2 past keys anymore
  ]


  @require_tf
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[str]			):






   """simple docstring"""






   for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
    self._test_export(UpperCamelCase_    ,		"""tf"""    ,		12    ,		**UpperCamelCase_			)


  @require_torch
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						Tuple			):






   """simple docstring"""






   for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
    self._test_export(UpperCamelCase_    ,		"""pt"""    ,		12    ,		**UpperCamelCase_			)


  @require_torch
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						Dict			):






   """simple docstring"""






   from transformers import BertModel

   UpperCAmelCase__      				=     ["""[UNK]""", """[SEP]""", """[CLS]""", """[PAD]""", """[MASK]""", """some""", """other""", """words"""]
   with NamedTemporaryFile(mode="""w+t"""			) as vocab_file:
    vocab_file.write("""\n""".join(UpperCamelCase_			)			)
    vocab_file.flush()
    UpperCAmelCase__      				=     BertTokenizerFast(vocab_file.name			)

   with TemporaryDirectory() as bert_save_dir:
    UpperCAmelCase__      				=     BertModel(BertConfig(vocab_size=len(UpperCamelCase_			)			)			)
    model.save_pretrained(UpperCamelCase_			)
    self._test_export(UpperCamelCase_    ,		"""pt"""    ,		12    ,		UpperCamelCase_			)


  @require_tf
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						Optional[Any]			):






   """simple docstring"""






   for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
    UpperCAmelCase__      				=     self._test_export(UpperCamelCase_    ,		"""tf"""    ,		12    ,		**UpperCamelCase_			)
    UpperCAmelCase__      				=     quantize(Path(UpperCamelCase_			)			)

    # Ensure the actual quantized model is not bigger than the original one
    if quantized_path.stat().st_size >= Path(UpperCamelCase_			).stat().st_size:
     self.fail("""Quantized model is bigger than initial ONNX model"""			)


  @require_torch
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						int			):






   """simple docstring"""






   for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
    UpperCAmelCase__      				=     self._test_export(UpperCamelCase_    ,		"""pt"""    ,		12    ,		**UpperCamelCase_			)
    UpperCAmelCase__      				=     quantize(UpperCamelCase_			)

    # Ensure the actual quantized model is not bigger than the original one
    if quantized_path.stat().st_size >= Path(UpperCamelCase_			).stat().st_size:
     self.fail("""Quantized model is bigger than initial ONNX model"""			)


  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[Any]    ,		_UpperCAmelCase     :						Optional[Any]    ,		_UpperCAmelCase     :						int    ,		_UpperCAmelCase     :						List[str]    ,		_UpperCAmelCase     :						int=None    ,		**_UpperCAmelCase     :						Any			):






   """simple docstring"""






   try:
    # Compute path
    with TemporaryDirectory() as tempdir:
     UpperCAmelCase__      				=     Path(UpperCamelCase_			).joinpath("""model.onnx"""			)

    # Remove folder if exists
    if path.parent.exists():
     path.parent.rmdir()

    # Export
    convert(UpperCamelCase_    ,		UpperCamelCase_    ,		UpperCamelCase_    ,		UpperCamelCase_    ,		UpperCamelCase_    ,		**UpperCamelCase_			)

    return path
   except Exception as e:
    self.fail(UpperCamelCase_			)


  @require_torch
  @require_tokenizers
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						str			):






   """simple docstring"""






   from transformers import BertModel

   UpperCAmelCase__      				=     BertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""			)			)
   UpperCAmelCase__      				=     BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""			)
   self._test_infer_dynamic_axis(UpperCamelCase_    ,		UpperCamelCase_    ,		"""pt"""			)


  @require_tf
  @require_tokenizers
  @slow
  def 				SCREAMING_SNAKE_CASE__      (    self     :						str			):






   """simple docstring"""






   from transformers import TFBertModel

   UpperCAmelCase__      				=     TFBertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random"""			)			)
   UpperCAmelCase__      				=     BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random"""			)
   self._test_infer_dynamic_axis(UpperCamelCase_    ,		UpperCamelCase_    ,		"""tf"""			)


  def 				SCREAMING_SNAKE_CASE__      (    self     :						Any    ,		_UpperCAmelCase     :						int    ,		_UpperCAmelCase     :						Optional[Any]    ,		_UpperCAmelCase     :						Any			):






   """simple docstring"""






   UpperCAmelCase__      				=     FeatureExtractionPipeline(UpperCamelCase_    ,		UpperCamelCase_			)

   UpperCAmelCase__      				=     ["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
   UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__				,      UpperCAmelCase__      				=     infer_shapes(UpperCamelCase_    ,		UpperCamelCase_			)

   # Assert all variables are present
   self.assertEqual(len(UpperCamelCase_			)    ,		len(UpperCamelCase_			)			)
   self.assertTrue(all(var_name in shapes for var_name in variable_names			)			)
   self.assertSequenceEqual(variable_names[:3]    ,		UpperCamelCase_			)
   self.assertSequenceEqual(variable_names[3:]    ,		UpperCamelCase_			)

   # Assert inputs are {0: batch, 1: sequence}
   for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
    self.assertDictEqual(shapes[var_name]    ,		{0: """batch""", 1: """sequence"""}			)

   # Assert outputs are {0: batch, 1: sequence} and {0: batch}
   self.assertDictEqual(shapes["""output_0"""]    ,		{0: """batch""", 1: """sequence"""}			)
   self.assertDictEqual(shapes["""output_1"""]    ,		{0: """batch"""}			)


  def 				SCREAMING_SNAKE_CASE__      (    self     :						List[str]			):






   """simple docstring"""






   UpperCAmelCase__      				=     ["""input_ids""", """attention_mask""", """token_type_ids"""]
   UpperCAmelCase__      				=     {"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
   UpperCAmelCase__				,      UpperCAmelCase__      				=     ensure_valid_input(FuncContiguousArgs()    ,		UpperCamelCase_    ,		UpperCamelCase_			)

   # Should have exactly the same number of args (all are valid)
   self.assertEqual(len(UpperCamelCase_			)    ,		3			)

   # Should have exactly the same input names
   self.assertEqual(set(UpperCamelCase_			)    ,		set(UpperCamelCase_			)			)

   # Parameter should be reordered according to their respective place in the function:
   # (input_ids, token_type_ids, attention_mask)
   self.assertEqual(UpperCamelCase_    ,		(tokens["""input_ids"""], tokens["""token_type_ids"""], tokens["""attention_mask"""])			)

   # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
   UpperCAmelCase__				,      UpperCAmelCase__      				=     ensure_valid_input(FuncNonContiguousArgs()    ,		UpperCamelCase_    ,		UpperCamelCase_			)

   # Should have exactly the one arg (all before the one not provided "some_other_args")
   self.assertEqual(len(UpperCamelCase_			)    ,		1			)
   self.assertEqual(len(UpperCamelCase_			)    ,		1			)

   # Should have only "input_ids"
   self.assertEqual(inputs_args[0]    ,		tokens["""input_ids"""]			)
   self.assertEqual(ordered_input_names[0]    ,		"""input_ids"""			)



  def 				SCREAMING_SNAKE_CASE__      (    self     :						Optional[Any]			):






   """simple docstring"""






   UpperCAmelCase__      				=     generate_identified_filename(Path("""/home/something/my_fake_model.onnx"""			)    ,		"""-test"""			)
   self.assertEqual("""/home/something/my_fake_model-test.onnx"""    ,		generated.as_posix()			)








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"""					)




