



'''simple docstring'''


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 	UpperCAmelCase	(		_UpperCamelCase):


 '''simple docstring'''




 def       lowercase_  (   self) ->      Optional[Any]:



      """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(__a							,  exist_ok=__a)
      a_       		=os.path.join(__a							,  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(__a							,  exist_ok=__a)




 def       lowercase_  (   self) ->      RealmTokenizer:



      """simple docstring"""







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




 def       lowercase_  (   self) ->      Dict:



      """simple docstring"""







      shutil.rmtree(self.tmpdirname)




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



      """simple docstring"""







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




 def       lowercase_  (   self) ->      int:



      """simple docstring"""







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




 def       lowercase_  (   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=__a							,  )
      return block_records




 def       lowercase_  (   self) ->      Optional[Any]:



      """simple docstring"""







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




 def       lowercase_  (   self) ->      Tuple:



      """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=__a							,  return_token_type_ids=__a							,  return_attention_mask=__a							,  ).input_ids
      a_       		=config.reader_seq_len

      a_       		=retriever(
          __a							,  __a							,  answer_ids=__a							,  max_length=__a							,  return_tensors="np")

      self.assertEqual(len(__a)							,  2)
      self.assertEqual(len(__a)							,  2)
      self.assertEqual(len(__a)							,  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) ->      List[Any]:



      """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=__a							,  return_token_type_ids=__a							,  return_attention_mask=__a							,  ).input_ids
      a_       		=config.reader_seq_len

      a_       		=retriever(
          __a							,  __a							,  answer_ids=__a							,  max_length=__a							,  return_tensors="np")

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




 def       lowercase_  (   self) ->      List[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")







'''simple docstring'''


from .testing import (
    are_the_same_tensors,
    execute_subprocess_async,
    require_bnb,
    require_cpu,
    require_cuda,
    require_huggingface_suite,
    require_mps,
    require_multi_gpu,
    require_multi_xpu,
    require_safetensors,
    require_single_gpu,
    require_single_xpu,
    require_torch_min_version,
    require_tpu,
    require_xpu,
    skip,
    slow,
)
from .training import RegressionDataset, RegressionModel, RegressionModelaXPU


from .scripts import test_script, test_sync, test_ops  # isort: skip


