



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 							_a							(					_a ):

			def 			__snake_case    (self )    ->					str:
						UpperCAmelCase_:						str					=			tempfile.mkdtemp()
						UpperCAmelCase_:						Tuple					=			5

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

			def 			__snake_case    (self )    ->					Optional[int]:
						return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname,					"""realm_tokenizer""" ) )

			def 			__snake_case    (self )    ->					Optional[int]:
						shutil.rmtree(self.tmpdirname )

			def 			__snake_case    (self )    ->					List[str]:
						UpperCAmelCase_:						List[Any]					=			RealmConfig(num_block_records=self.num_block_records )
						return config

			def 			__snake_case    (self )    ->					Any:
						UpperCAmelCase_:						str					=			Dataset.from_dict(
						    {
						        """id""": ["""0""", """1"""],
						        """question""": ["""foo""", """bar"""],
						        """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
						    } )
						return dataset

			def 			__snake_case    (self )    ->					Dict:
						UpperCAmelCase_:						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=snake_case_,					)
						return block_records

			def 			__snake_case    (self )    ->					List[str]:
						UpperCAmelCase_:						Dict					=			RealmRetriever(
						    block_records=self.get_dummy_block_records(),					tokenizer=self.get_tokenizer(),					)
						return retriever

			def 			__snake_case    (self )    ->					Any:
						UpperCAmelCase_:						Tuple					=			self.get_config()
						UpperCAmelCase_:						Any					=			self.get_dummy_retriever()
						UpperCAmelCase_:						str					=			retriever.tokenizer

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

						UpperCAmelCase_:						Optional[int]					=			retriever(
						    snake_case_,					snake_case_,					answer_ids=snake_case_,					max_length=snake_case_,					return_tensors="""np""" )

						self.assertEqual(len(snake_case_ ),					2 )
						self.assertEqual(len(snake_case_ ),					2 )
						self.assertEqual(len(snake_case_ ),					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 			__snake_case    (self )    ->					Optional[Any]:
						UpperCAmelCase_:						List[str]					=			self.get_config()
						UpperCAmelCase_:						List[str]					=			self.get_dummy_retriever()
						UpperCAmelCase_:						str					=			retriever.tokenizer

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

						UpperCAmelCase_:						List[str]					=			retriever(
						    snake_case_,					snake_case_,					answer_ids=snake_case_,					max_length=snake_case_,					return_tensors="""np""" )

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

			def 			__snake_case    (self )    ->					List[Any]:
						UpperCAmelCase_:						List[str]					=			self.get_dummy_retriever()
						retriever.save_pretrained(os.path.join(self.tmpdirname,					"""realm_block_records""" ) )

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

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











'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__a          =		{
    "configuration_bloom": ["BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP", "BloomConfig", "BloomOnnxConfig"],
}
try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __a          =		["BloomTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __a          =		[
           "BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST",
           "BloomForCausalLM",
           "BloomModel",
           "BloomPreTrainedModel",
           "BloomForSequenceClassification",
           "BloomForTokenClassification",
           "BloomForQuestionAnswering",
       ]

if TYPE_CHECKING:
       from .configuration_bloom import BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP, BloomConfig, BloomOnnxConfig

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tokenization_bloom_fast import BloomTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_bloom import (
                  BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST,
                  BloomForCausalLM,
                  BloomForQuestionAnswering,
                  BloomForSequenceClassification,
                  BloomForTokenClassification,
                  BloomModel,
                  BloomPreTrainedModel,
              )

else:
       import sys

       __a          =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
