





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

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   List[Any]:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   tempfile.mkdtemp()
       UpperCamelCase__							:       str      =   5

       # Realm tok
       UpperCamelCase__							:       Tuple      =   [
           '''[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__							:       str      =   os.path.join(self.tmpdirname,	'''realm_tokenizer'''							)
       os.makedirs(__magic_name__,	exist_ok=__magic_name__							)
       UpperCamelCase__							:       List[str]      =   os.path.join(__magic_name__,	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(__magic_name__,	exist_ok=__magic_name__							)




   def    UpperCamelCase__						(					self							)				->   RealmTokenizer:






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




   def    UpperCamelCase__						(					self							)				->   List[str]:






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




   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       List[Any]      =   RealmConfig(num_block_records=self.num_block_records							)
       return config




   def    UpperCamelCase__						(					self							)				->   List[str]:






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




   def    UpperCamelCase__						(					self							)				->   Union[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       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=__magic_name__,	)
       return block_records




   def    UpperCamelCase__						(					self							)				->   List[Any]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   RealmRetriever(
           block_records=self.get_dummy_block_records(),	tokenizer=self.get_tokenizer(),	)
       return retriever




   def    UpperCamelCase__						(					self							)				->   Union[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   self.get_config()
       UpperCamelCase__							:       int      =   self.get_dummy_retriever()
       UpperCamelCase__							:       int      =   retriever.tokenizer

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

       UpperCamelCase__							:       str      =   retriever(
           __magic_name__,	__magic_name__,	answer_ids=__magic_name__,	max_length=__magic_name__,	return_tensors='''np'''							)

       self.assertEqual(len(__magic_name__							),	2							)
       self.assertEqual(len(__magic_name__							),	2							)
       self.assertEqual(len(__magic_name__							),	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							)				->   List[str]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   self.get_config()
       UpperCamelCase__							:       Dict      =   self.get_dummy_retriever()
       UpperCamelCase__							:       Optional[Any]      =   retriever.tokenizer

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

       UpperCamelCase__							:       Tuple      =   retriever(
           __magic_name__,	__magic_name__,	answer_ids=__magic_name__,	max_length=__magic_name__,	return_tensors='''np'''							)

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




   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   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__							:       str      =   os.path.join(
               os.path.join(self.tmpdirname,	'''realm_block_records'''							),	_REALM_BLOCK_RECORDS_FILENAME							)
           UpperCamelCase__							:       Any      =   RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''							)

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








from manim import *






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       UpperCamelCase__							:       Any      =   Rectangle(height=0.5,	width=0.5							)
       UpperCamelCase__							:       Any      =   Rectangle(height=0.25,	width=0.25							)
       UpperCamelCase__							:       Dict      =   Rectangle(height=0.46,	width=0.46							).set_stroke(width=0							)

       UpperCamelCase__							:       List[Any]      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Optional[int]      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       str      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       List[str]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Tuple      =   VGroup(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       List[str]      =   Text('''CPU''',	font_size=24							)
       UpperCamelCase__							:       Any      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       cpu.move_to([-2.5, -0.5, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       int      =   [mem.copy() for i in range(4							)]
       UpperCamelCase__							:       str      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Tuple      =   Text('''GPU''',	font_size=24							)
       UpperCamelCase__							:       str      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       gpu.move_to([-1, -1, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       str      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Tuple      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)

       UpperCamelCase__							:       Tuple      =   Text('''Model''',	font_size=24							)
       UpperCamelCase__							:       Tuple      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       model.move_to([3, -1.0, 0]							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       Union[str, Any]      =   []
       UpperCamelCase__							:       Any      =   []
       UpperCamelCase__							:       List[Any]      =   []

       for i, rect in enumerate(__magic_name__							):
           rect.set_stroke(__magic_name__							)

           UpperCamelCase__							:       Tuple      =   Rectangle(height=0.46 / 4,	width=0.46 / 3							).set_stroke(width=0.0							).set_fill(__magic_name__,	opacity=0.7							)

           if i == 0:
               cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT							),	buff=0.02,	direction=__magic_name__							)
               cpu_target.set_x(cpu_target.get_x() + 0.1							)
           elif i == 3:
               cpu_target.next_to(model_cpu_arr[0],	direction=__magic_name__,	buff=0.0							)
           else:
               cpu_target.next_to(model_cpu_arr[i - 1],	direction=__magic_name__,	buff=0.0							)
           self.add(__magic_name__							)
           model_cpu_arr.append(__magic_name__							)

       self.add(*__magic_name__,	*__magic_name__,	*__magic_name__							)

       UpperCamelCase__							:       Tuple      =   [mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Dict      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)

       UpperCamelCase__							:       List[str]      =   Text('''Loaded Checkpoint''',	font_size=24							)
       UpperCamelCase__							:       Optional[Any]      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       checkpoint.move_to([3, 0.5, 0]							)

       self.add(__magic_name__							)

       UpperCamelCase__							:       Tuple      =   []
       UpperCamelCase__							:       List[str]      =   []

       for i, rect in enumerate(__magic_name__							):
           UpperCamelCase__							:       Optional[int]      =   fill.copy().set_fill(__magic_name__,	opacity=0.7							)
           target.move_to(__magic_name__							)
           ckpt_arr.append(__magic_name__							)

           UpperCamelCase__							:       int      =   target.copy()
           if i < 5:
               cpu_target.move_to(cpu_left_col_base[i + 1]							)
           else:
               cpu_target.move_to(cpu_right_col_base[i - 5]							)
           ckpt_cpu_arr.append(__magic_name__							)
       self.add(*__magic_name__,	*__magic_name__							)

       UpperCamelCase__							:       List[str]      =   Square(side_length=2.2							)
       key.move_to([-5, 2, 0]							)

       UpperCamelCase__							:       List[str]      =   MarkupText(
           f"<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model",	font_size=18,	)

       key_text.move_to([-5, 2.4, 0]							)

       self.add(__magic_name__,	__magic_name__							)

       UpperCamelCase__							:       Any      =   MarkupText(
           f"<span fgcolor='{BLUE}'>●</span> Checkpoint",	font_size=18,	)

       blue_text.next_to(__magic_name__,	DOWN * 2.4,	aligned_edge=key_text.get_left()							)
       self.add(__magic_name__							)

       UpperCamelCase__							:       Dict      =   MarkupText(
           f"Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.",	font_size=24,	)
       step_a.move_to([2, 2, 0]							)

       UpperCamelCase__							:       Any      =   [meta_mem.copy() for i in range(6							)]
       UpperCamelCase__							:       int      =   [meta_mem.copy() for i in range(6							)]
       UpperCamelCase__							:       Optional[int]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Union[str, Any]      =   VGroup(*__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Union[str, Any]      =   VGroup(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0							)
       UpperCamelCase__							:       Any      =   Text('''Disk''',	font_size=24							)
       UpperCamelCase__							:       List[Any]      =   Group(__magic_name__,	__magic_name__							).arrange(__magic_name__,	buff=0.5,	aligned_edge=__magic_name__							)
       disk.move_to([-4.0, -1.25, 0]							)
       self.play(Write(__magic_name__,	run_time=3							),	Write(__magic_name__,	run_time=1							),	Create(__magic_name__,	run_time=1							)							)

       UpperCamelCase__							:       Union[str, Any]      =   []
       for i, rect in enumerate(__magic_name__							):
           UpperCamelCase__							:       Union[str, Any]      =   rect.copy()
           target.generate_target()
           target.target.move_to(disk_left_col_base[i]							).scale(0.5							)
           animations.append(MoveToTarget(__magic_name__,	run_time=1.5							)							)
       self.play(*__magic_name__							)

       self.play(FadeOut(__magic_name__							)							)

       UpperCamelCase__							:       Optional[int]      =   MarkupText(f"Then, the checkpoint is removed from memory\nthrough garbage collection.",	font_size=24							)
       step_a.move_to([2, 2, 0]							)

       self.play(Write(__magic_name__,	run_time=3							)							)

       self.play(
           FadeOut(__magic_name__,	__magic_name__,	*__magic_name__,	*__magic_name__							),	)

       self.wait()

