



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




       """simple docstring"""






       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		tempfile.mkdtemp()
             UpperCAmelCase_  :	List[str]				      =		5

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



       def  A_   (     self):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             shutil.rmtree(self.tmpdirname)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		RealmConfig(num_block_records=self.num_block_records)
             return config



       def  A_   (     self):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             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=lowercase     ,)
             return block_records



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		RealmRetriever(
                 block_records=self.get_dummy_block_records()     ,tokenizer=self.get_tokenizer()     ,)
             return retriever



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		self.get_config()
             UpperCAmelCase_  :	Union[str, Any]				      =		self.get_dummy_retriever()
             UpperCAmelCase_  :	Tuple				      =		retriever.tokenizer

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

             UpperCAmelCase_      ,      UpperCAmelCase_      ,      UpperCAmelCase_      ,      UpperCAmelCase_  :	Dict				      =		retriever(
                 lowercase     ,lowercase     ,answer_ids=lowercase     ,max_length=lowercase     ,return_tensors="np")

             self.assertEqual(len(lowercase)     ,2)
             self.assertEqual(len(lowercase)     ,2)
             self.assertEqual(len(lowercase)     ,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  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[str]				      =		self.get_config()
             UpperCAmelCase_  :	List[str]				      =		self.get_dummy_retriever()
             UpperCAmelCase_  :	Optional[int]				      =		retriever.tokenizer

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

             UpperCAmelCase_      ,      UpperCAmelCase_      ,      UpperCAmelCase_      ,      UpperCAmelCase_  :	Optional[Any]				      =		retriever(
                 lowercase     ,lowercase     ,answer_ids=lowercase     ,max_length=lowercase     ,return_tensors="np")

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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[int]				      =		self.get_dummy_retriever()
             retriever.save_pretrained(os.path.join(self.tmpdirname     ,"realm_block_records"))

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

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





import logging

from transformers import PretrainedConfig


__lowerCamelCase         =			logging.getLogger(__name__)


__lowerCamelCase         =			{
    '''bertabs-finetuned-cnndm''': '''https://huggingface.co/remi/bertabs-finetuned-cnndm-extractive-abstractive-summarization/resolve/main/config.json''',
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""bertabs"""



       def __init__(     self     ,lowercase=30522     ,lowercase=512     ,lowercase=6     ,lowercase=512     ,lowercase=8     ,lowercase=512     ,lowercase=0.2     ,lowercase=6     ,lowercase=768     ,lowercase=8     ,lowercase=2048     ,lowercase=0.2     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(**lowercase)

             UpperCAmelCase_  :	Any				      =		vocab_size
             UpperCAmelCase_  :	int				      =		max_pos

             UpperCAmelCase_  :	List[str]				      =		enc_layers
             UpperCAmelCase_  :	List[str]				      =		enc_hidden_size
             UpperCAmelCase_  :	Any				      =		enc_heads
             UpperCAmelCase_  :	Any				      =		enc_ff_size
             UpperCAmelCase_  :	Optional[int]				      =		enc_dropout

             UpperCAmelCase_  :	List[str]				      =		dec_layers
             UpperCAmelCase_  :	List[Any]				      =		dec_hidden_size
             UpperCAmelCase_  :	List[str]				      =		dec_heads
             UpperCAmelCase_  :	List[Any]				      =		dec_ff_size
             UpperCAmelCase_  :	List[str]				      =		dec_dropout
