


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





							"""simple docstring"""







							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					tempfile.mkdtemp()
										SCREAMING_SNAKE_CASE_		:     str					=					5

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

										SCREAMING_SNAKE_CASE_		:     str					=					os.path.join(self.tmpdirname			,						'realm_block_records'       )
										os.makedirs(_lowerCamelCase			,						exist_ok=_lowerCamelCase       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										shutil.rmtree(self.tmpdirname       )






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					RealmConfig(num_block_records=self.num_block_records       )
										return config






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					Dataset.from_dict(
										    {
										        'id': ['0', '1'],
										        'question': ['foo', 'bar'],
										        'answers': [['Foo', 'Bar'], ['Bar']],
										    }       )
										return dataset






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     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=_lowerCamelCase			,						)
										return block_records






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     Tuple					=					retriever.tokenizer

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

										SCREAMING_SNAKE_CASE_		:     Any					=					retriever(
										    _lowerCamelCase			,						_lowerCamelCase			,						answer_ids=_lowerCamelCase			,						max_length=_lowerCamelCase			,						return_tensors='np'       )

										self.assertEqual(len(_lowerCamelCase       )			,						2       )
										self.assertEqual(len(_lowerCamelCase       )			,						2       )
										self.assertEqual(len(_lowerCamelCase       )			,						2       )
										self.assertEqual(concat_inputs.input_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.attention_mask.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.token_type_ids.shape			,						(2, 1_0)       )
										self.assertEqual(concat_inputs.special_tokens_mask.shape			,						(2, 1_0)       )
										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       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.get_config()
										SCREAMING_SNAKE_CASE_		:     int					=					self.get_dummy_retriever()
										SCREAMING_SNAKE_CASE_		:     Any					=					retriever.tokenizer

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

										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					retriever(
										    _lowerCamelCase			,						_lowerCamelCase			,						answer_ids=_lowerCamelCase			,						max_length=_lowerCamelCase			,						return_tensors='np'       )

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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					self.get_dummy_retriever()
										retriever.save_pretrained(os.path.join(self.tmpdirname			,						'realm_block_records'       )       )

										# Test local path
										SCREAMING_SNAKE_CASE_		:     int					=					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:
													SCREAMING_SNAKE_CASE_		:     Dict					=					os.path.join(
													    os.path.join(self.tmpdirname			,						'realm_block_records'       )			,						_REALM_BLOCK_RECORDS_FILENAME       )
													SCREAMING_SNAKE_CASE_		:     int					=					RealmRetriever.from_pretrained('google/realm-cc-news-pretrained-openqa'       )

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







"""simple docstring"""






# We ignore warnings about stepping the scheduler since we step it ourselves during gradient accumulation

import warnings

from .state import AcceleratorState, GradientState


warnings.filterwarnings('ignore', category=UserWarning, module='torch.optim.lr_scheduler')





class 				a__					:



     def __init__(							self :List[str]      ,  _lowerCamelCase :Any      ,  _lowerCamelCase :Optional[Any]      ,  _lowerCamelCase :bool = True      ,  _lowerCamelCase :bool = False   ):







      '''simple docstring'''



      UpperCamelCase_     :					List[Any]   	=scheduler
      UpperCamelCase_     :					Any   	=optimizers if isinstance(_lowerCamelCase      ,  (list, tuple)   ) else [optimizers]
      UpperCamelCase_     :					Dict   	=split_batches
      UpperCamelCase_     :					str   	=step_with_optimizer
      UpperCamelCase_     :					Dict   	=GradientState()



     def  lowerCamelCase_						(							self :Optional[int]      ,  *_lowerCamelCase :str      ,  **_lowerCamelCase :Union[str, Any]   ):







      '''simple docstring'''



      if not self.step_with_optimizer:
       # No link between scheduler and optimizer -> just step
       self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
       return

      # Otherwise, first make sure the optimizer was stepped.
      if not self.gradient_state.sync_gradients:
       if self.gradient_state.adjust_scheduler:
        self.scheduler._step_count += 1
       return

      for opt in self.optimizers:
       if opt.step_was_skipped:
        return
      if self.split_batches:
       # Split batches -> the training dataloader batch size is not changed so one step per training step
       self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
      else:
       # Otherwise the training dataloader batch size was multiplied by `num_processes`, so we need to do
       # num_processes steps per training step
       UpperCamelCase_     :					Optional[Any]   	=AcceleratorState().num_processes
       for _ in range(_lowerCamelCase   ):
        # Special case when using OneCycle and `drop_last` was not used
        if hasattr(self.scheduler      ,  'total_steps'   ):
         if self.scheduler._step_count <= self.scheduler.total_steps:
          self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )
        else:
         self.scheduler.step(*_lowerCamelCase      ,  **_lowerCamelCase   )



     def  lowerCamelCase_						(							self :Dict   ):







      '''simple docstring'''



      return self.scheduler.get_last_lr()



     def  lowerCamelCase_						(							self :Tuple   ):







      '''simple docstring'''



      return self.scheduler.state_dict()



     def  lowerCamelCase_						(							self :Any      ,  _lowerCamelCase :Any   ):







      '''simple docstring'''



      self.scheduler.load_state_dict(_lowerCamelCase   )



     def  lowerCamelCase_						(							self :Any   ):







      '''simple docstring'''



      return self.scheduler.get_lr()




     def  lowerCamelCase_						(							self :Optional[int]      ,  *_lowerCamelCase :Union[str, Any]      ,  **_lowerCamelCase :Optional[int]   ):







      '''simple docstring'''



      return self.scheduler.print_lr(*_lowerCamelCase      ,  **_lowerCamelCase   )
