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






   """simple docstring"""
   def 		_lowerCamelCase (self				)    ->     Any:
        lowercase_					:       Optional[int]     				=					tempfile.mkdtemp()
        lowercase_					:       str     				=					5

        # Realm tok
        lowercase_					:       Optional[int]     				=					[
            '[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_					:       Optional[Any]     				=					os.path.join(self.tmpdirname       ,				'realm_tokenizer'				)
        os.makedirs(_a       ,				exist_ok=_a				)
        lowercase_					:       str     				=					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_					:       Optional[Any]     				=					os.path.join(self.tmpdirname       ,				'realm_block_records'				)
        os.makedirs(_a       ,				exist_ok=_a				)
   def 		_lowerCamelCase (self				)    ->     RealmTokenizer:
        return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,				'realm_tokenizer'				)				)
   def 		_lowerCamelCase (self				)    ->     Dict:
        shutil.rmtree(self.tmpdirname				)
   def 		_lowerCamelCase (self				)    ->     str:
        lowercase_					:       str     				=					RealmConfig(num_block_records=self.num_block_records				)
        return config
   def 		_lowerCamelCase (self				)    ->     Any:
        lowercase_					:       Union[str, Any]     				=					Dataset.from_dict(
            {
                'id': ['0', '1'],
                'question': ['foo', 'bar'],
                'answers': [['Foo', 'Bar'], ['Bar']],
            }				)
        return dataset
   def 		_lowerCamelCase (self				)    ->     Optional[Any]:
        lowercase_					:       Dict     				=					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]:
        lowercase_					:       int     				=					RealmRetriever(
            block_records=self.get_dummy_block_records()       ,				tokenizer=self.get_tokenizer()       ,				)
        return retriever
   def 		_lowerCamelCase (self				)    ->     str:
        lowercase_					:       int     				=					self.get_config()
        lowercase_					:       Tuple     				=					self.get_dummy_retriever()
        lowercase_					:       Union[str, Any]     				=					retriever.tokenizer

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

        lowercase_				,lowercase_				,lowercase_				,lowercase_					:       Optional[int]     				=					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, 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				)    ->     List[Any]:
        lowercase_					:       int     				=					self.get_config()
        lowercase_					:       Optional[int]     				=					self.get_dummy_retriever()
        lowercase_					:       Union[str, Any]     				=					retriever.tokenizer

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

        lowercase_				,lowercase_				,lowercase_				,lowercase_					:       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				)    ->     int:
        lowercase_					:       Optional[int]     				=					self.get_dummy_retriever()
        retriever.save_pretrained(os.path.join(self.tmpdirname       ,				'realm_block_records'				)				)

        # Test local path
        lowercase_					:       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:
             lowercase_					:       Optional[int]     				=					os.path.join(
                 os.path.join(self.tmpdirname       ,				'realm_block_records'				)       ,				_REALM_BLOCK_RECORDS_FILENAME				)
             lowercase_					:       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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_A								=   {
    'configuration_bigbird_pegasus': [
        'BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'BigBirdPegasusConfig',
        'BigBirdPegasusOnnxConfig',
    ],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _A								=   [
     'BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST',
     'BigBirdPegasusForCausalLM',
     'BigBirdPegasusForConditionalGeneration',
     'BigBirdPegasusForQuestionAnswering',
     'BigBirdPegasusForSequenceClassification',
     'BigBirdPegasusModel',
     'BigBirdPegasusPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_bigbird_pegasus import (
     BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
     BigBirdPegasusConfig,
     BigBirdPegasusOnnxConfig,
 )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bigbird_pegasus import (
      BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
      BigBirdPegasusForCausalLM,
      BigBirdPegasusForConditionalGeneration,
      BigBirdPegasusForQuestionAnswering,
      BigBirdPegasusForSequenceClassification,
      BigBirdPegasusModel,
      BigBirdPegasusPreTrainedModel,
  )


else:
 import sys

 _A								=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




