






'''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						(		snake_case			):




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     tempfile.mkdtemp()
       lowercase_	:    Optional[int]           =     5

       # Realm tok
       lowercase_	:    Tuple           =     [
           "[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",
       ]
       lowercase_	:    int           =     os.path.join(self.tmpdirname  ,     "realm_tokenizer"			)
       os.makedirs(a  ,     exist_ok=a			)
       lowercase_	:    Union[str, Any]           =     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]			)			)

       lowercase_	:    List[Any]           =     os.path.join(self.tmpdirname  ,     "realm_block_records"			)
       os.makedirs(a  ,     exist_ok=a			)




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






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




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






       shutil.rmtree(self.tmpdirname			)




    def 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    int           =     RealmConfig(num_block_records=self.num_block_records			)
       return config




    def 				lowerCAmelCase__				(			self  :       List[str]			):






       '''simple docstring'''






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




    def 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    List[str]           =     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 				lowerCAmelCase__				(			self  :       List[Any]			):






       '''simple docstring'''






       lowercase_	:    Optional[Any]           =     RealmRetriever(
           block_records=self.get_dummy_block_records()  ,     tokenizer=self.get_tokenizer()  ,     )
       return retriever




    def 				lowerCAmelCase__				(			self  :       Any			):






       '''simple docstring'''






       lowercase_	:    Dict           =     self.get_config()
       lowercase_	:    Any           =     self.get_dummy_retriever()
       lowercase_	:    Tuple           =     retriever.tokenizer

       lowercase_	:    str           =     np.array([0, 3]  ,     dtype="long"			)
       lowercase_	:    int           =     tokenizer(["Test question"]			).input_ids
       lowercase_	:    List[Any]           =     tokenizer(
           ["the fourth"]  ,     add_special_tokens=a  ,     return_token_type_ids=a  ,     return_attention_mask=a  ,     ).input_ids
       lowercase_	:    Union[str, Any]           =     config.reader_seq_len

       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    Optional[Any]           =     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 				lowerCAmelCase__				(			self  :       int			):






       '''simple docstring'''






       lowercase_	:    Tuple           =     self.get_config()
       lowercase_	:    Optional[int]           =     self.get_dummy_retriever()
       lowercase_	:    int           =     retriever.tokenizer

       lowercase_	:    Tuple           =     np.array([0, 3, 5]  ,     dtype="long"			)
       lowercase_	:    List[str]           =     tokenizer(["Test question"]			).input_ids
       lowercase_	:    Dict           =     tokenizer(
           ["the fourth", "longer longer"]  ,     add_special_tokens=a  ,     return_token_type_ids=a  ,     return_attention_mask=a  ,     ).input_ids
       lowercase_	:    Union[str, Any]           =     config.reader_seq_len

       lowercase_    ,		lowercase_    ,		lowercase_    ,		lowercase_	:    Union[str, Any]           =     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 				lowerCAmelCase__				(			self  :       Optional[Any]			):






       '''simple docstring'''






       lowercase_	:    List[Any]           =     self.get_dummy_retriever()
       retriever.save_pretrained(os.path.join(self.tmpdirname  ,     "realm_block_records"			)			)

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

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












'''simple docstring'''
import os
from collections import namedtuple

import pytest

from datasets import ClassLabel, Features, Sequence, Value
from datasets.commands.test import TestCommand
from datasets.info import DatasetInfo, DatasetInfosDict


UpperCamelCase__ =			namedtuple(
    '_TestCommandArgs',
    [
        'dataset',
        'name',
        'cache_dir',
        'data_dir',
        'all_configs',
        'save_infos',
        'ignore_verifications',
        'force_redownload',
        'clear_cache',
    ],
    defaults=[None, None, None, False, False, False, False, False],
)




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase							):


   """simple docstring"""




   return (abs(source - target							) / target) < 0.01





@pytest.mark.integration
def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase							):


   """simple docstring"""




   lowercase_	:    Optional[int]           =     _TestCommandArgs(dataset=_UpperCamelCase						,		all_configs=_UpperCamelCase						,		save_infos=_UpperCamelCase							)
   lowercase_	:    int           =     TestCommand(*_UpperCamelCase							)
   test_command.run()
   lowercase_	:    List[str]           =     os.path.join(_UpperCamelCase						,		"README.md"							)
   assert os.path.exists(_UpperCamelCase							)
   lowercase_	:    Any           =     DatasetInfosDict.from_directory(_UpperCamelCase							)
   lowercase_	:    Optional[int]           =     DatasetInfosDict(
       {
           "default": DatasetInfo(
               features=Features(
                   {
                       "tokens": Sequence(Value("string"							)							),
                       "ner_tags": Sequence(
                           ClassLabel(names=["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]							)							),
                       "langs": Sequence(Value("string"							)							),
                       "spans": Sequence(Value("string"							)							),
                   }							)						,		splits=[
                   {
                       "name": "train",
                       "num_bytes": 235_1563,
                       "num_examples": 1_0000,
                   },
                   {
                       "name": "validation",
                       "num_bytes": 23_8418,
                       "num_examples": 1000,
                   },
               ]						,		download_size=394_0680						,		dataset_size=258_9981						,		)
       }							)
   assert dataset_infos.keys() == expected_dataset_infos.keys()
   for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
      lowercase_    ,		lowercase_	:    Optional[int]           =     getattr(dataset_infos["default"]						,		_UpperCamelCase							), getattr(expected_dataset_infos["default"]						,		_UpperCamelCase							)
      if key == "num_bytes":
         assert is_apercent_close(_UpperCamelCase						,		_UpperCamelCase							)
      elif key == "splits":
         assert list(_UpperCamelCase							) == list(_UpperCamelCase							)
         for split in result:
            assert result[split].name == expected[split].name
            assert result[split].num_examples == expected[split].num_examples
            assert is_apercent_close(result[split].num_bytes						,		expected[split].num_bytes							)
      else:
         result == expected




