






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




							def 	lowerCamelCase_       (							self:       int  ):
														lowerCamelCase__      :			Optional[int] 						= tempfile.mkdtemp()
														lowerCamelCase__      :			int 						= 5

														# Realm tok
														lowerCamelCase__      :			List[str] 						= [
														    """[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""",
														]
														lowerCamelCase__      :			List[str] 						= os.path.join(self.tmpdirname    ,							"""realm_tokenizer"""  )
														os.makedirs(UpperCamelCase__    ,							exist_ok=UpperCamelCase__  )
														lowerCamelCase__      :			Optional[Any] 						= 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]  )  )

														lowerCamelCase__      :			str 						= os.path.join(self.tmpdirname    ,							"""realm_block_records"""  )
														os.makedirs(UpperCamelCase__    ,							exist_ok=UpperCamelCase__  )




							def 	lowerCamelCase_       (							self:       Union[str, Any]  ):
														return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname    ,							"""realm_tokenizer"""  )  )




							def 	lowerCamelCase_       (							self:       Union[str, Any]  ):
														shutil.rmtree(self.tmpdirname  )




							def 	lowerCamelCase_       (							self:       Optional[int]  ):
														lowerCamelCase__      :			int 						= RealmConfig(num_block_records=self.num_block_records  )
														return config




							def 	lowerCamelCase_       (							self:       Any  ):
														lowerCamelCase__      :			Tuple 						= Dataset.from_dict(
														    {
														        """id""": ["""0""", """1"""],
														        """question""": ["""foo""", """bar"""],
														        """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
														    }  )
														return dataset




							def 	lowerCamelCase_       (							self:       int  ):
														lowerCamelCase__      :			Union[str, Any] 						= 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:       Any  ):
														lowerCamelCase__      :			str 						= RealmRetriever(
														    block_records=self.get_dummy_block_records()    ,							tokenizer=self.get_tokenizer()    ,							)
														return retriever




							def 	lowerCamelCase_       (							self:       Tuple  ):
														lowerCamelCase__      :			int 						= self.get_config()
														lowerCamelCase__      :			Tuple 						= self.get_dummy_retriever()
														lowerCamelCase__      :			Any 						= retriever.tokenizer

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

														lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__      :			Tuple 						= 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  ):
														lowerCamelCase__      :			Tuple 						= self.get_config()
														lowerCamelCase__      :			Tuple 						= self.get_dummy_retriever()
														lowerCamelCase__      :			List[Any] 						= retriever.tokenizer

														lowerCamelCase__      :			str 						= np.array([0, 3, 5]    ,							dtype="""long"""  )
														lowerCamelCase__      :			Dict 						= tokenizer(["""Test question"""]  ).input_ids
														lowerCamelCase__      :			str 						= tokenizer(
														    ["""the fourth""", """longer longer"""]    ,							add_special_tokens=UpperCamelCase__    ,							return_token_type_ids=UpperCamelCase__    ,							return_attention_mask=UpperCamelCase__    ,							).input_ids
														lowerCamelCase__      :			Dict 						= config.reader_seq_len

														lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__						,	lowerCamelCase__      :			Union[str, Any] 						= 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  ):
														lowerCamelCase__      :			int 						= self.get_dummy_retriever()
														retriever.save_pretrained(os.path.join(self.tmpdirname    ,							"""realm_block_records"""  )  )

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

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









'''simple docstring'''







_A			:    List[str]					  ='''0.21.0'''

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
			from .utils import rich

