



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





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       tempfile.mkdtemp()
							_UpperCamelCase :		Union[str, Any]			    =       5

							# Realm tok
							_UpperCamelCase :		Any			    =       [
							    "[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 :		Dict			    =       os.path.join(self.tmpdirname			,'realm_tokenizer'			)
							os.makedirs(lowerCamelCase__			,exist_ok=lowerCamelCase__			)
							_UpperCamelCase :		Optional[Any]			    =       os.path.join(lowerCamelCase__			,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 :		Tuple			    =       os.path.join(self.tmpdirname			,'realm_block_records'			)
							os.makedirs(lowerCamelCase__			,exist_ok=lowerCamelCase__			)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,'realm_tokenizer'			)			)





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							shutil.rmtree(self.tmpdirname			)





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[Any]			    =       RealmConfig(num_block_records=self.num_block_records			)
							return config





	def      UpperCamelCase_  (			self :    List[str]			):







							'''simple docstring'''
							_UpperCamelCase :		Any			    =       Dataset.from_dict(
							    {
							        'id': ['0', '1'],
							        'question': ['foo', 'bar'],
							        'answers': [['Foo', 'Bar'], ['Bar']],
							    }			)
							return dataset





	def      UpperCamelCase_  (			self :    str			):







							'''simple docstring'''
							_UpperCamelCase :		int			    =       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=lowerCamelCase__			,)
							return block_records





	def      UpperCamelCase_  (			self :    Optional[Any]			):







							'''simple docstring'''
							_UpperCamelCase :		List[str]			    =       RealmRetriever(
							    block_records=self.get_dummy_block_records()			,tokenizer=self.get_tokenizer()			,)
							return retriever





	def      UpperCamelCase_  (			self :    Tuple			):







							'''simple docstring'''
							_UpperCamelCase :		Optional[int]			    =       self.get_config()
							_UpperCamelCase :		Optional[int]			    =       self.get_dummy_retriever()
							_UpperCamelCase :		Any			    =       retriever.tokenizer

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

							_UpperCamelCase :		List[str]			    =       retriever(
							    lowerCamelCase__			,lowerCamelCase__			,answer_ids=lowerCamelCase__			,max_length=lowerCamelCase__			,return_tensors='np'			)

							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(len(lowerCamelCase__			)			,2			)
							self.assertEqual(len(lowerCamelCase__			)			,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      UpperCamelCase_  (			self :    Any			):







							'''simple docstring'''
							_UpperCamelCase :		Union[str, Any]			    =       self.get_config()
							_UpperCamelCase :		int			    =       self.get_dummy_retriever()
							_UpperCamelCase :		List[Any]			    =       retriever.tokenizer

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

							_UpperCamelCase :		int			    =       retriever(
							    lowerCamelCase__			,lowerCamelCase__			,answer_ids=lowerCamelCase__			,max_length=lowerCamelCase__			,return_tensors='np'			)

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





	def      UpperCamelCase_  (			self :    Dict			):







							'''simple docstring'''
							_UpperCamelCase :		str			    =       self.get_dummy_retriever()
							retriever.save_pretrained(os.path.join(self.tmpdirname			,'realm_block_records'			)			)

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

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







'''simple docstring'''
from torch import nn






def A__   (					UpperCAmelCase_							):
						if act_fn in ["swish", "silu"]:
												return nn.SiLU()
						elif act_fn == "mish":
												return nn.Mish()
						elif act_fn == "gelu":
												return nn.GELU()
						else:
												raise ValueError(f'Unsupported activation function: {act_fn}'							)


