



"""simple docstring"""
import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment

def       _snake_case						(						):
 UpperCAmelCase	:					List[Any]	=	torch.nn.Linear(2			,		4						)
 UpperCAmelCase	:					Optional[int]	=	torch.optim.AdamW(model.parameters()			,		lr=1.0						)
 UpperCAmelCase	:					int	=	torch.optim.lr_scheduler.OneCycleLR(UpperCamelCase			,		max_lr=0.01			,		steps_per_epoch=2			,		epochs=1						)
 UpperCAmelCase	:					List[str]	=	DataLoader(TensorDataset(torch.tensor([1, 2, 3]						)						)						)
 UpperCAmelCase	:					Tuple	=	DataLoader(TensorDataset(torch.tensor([4, 5, 6]						)						)						)

 return model, optimizer, scheduler, train_dl, valid_dl

def       _snake_case						(						UpperCamelCase						:		int						):
 return (model.weight.abs().sum() + model.bias.abs().sum()).item()

def       _snake_case						(						UpperCamelCase						:		List[Any]						):
 UpperCAmelCase	:					Tuple	=	torch.nn.Linear(*tuple(model.weight.T.shape						)						).state_dict()
 model.load_state_dict(UpperCamelCase						)



class       SCREAMING_SNAKE_CASE__  (					UpperCAmelCase__ ):

   @require_cuda
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	Accelerator()
    assert PartialState._shared_state["_cpu"] is False
    assert PartialState._shared_state["device"].type == "cuda"
    with self.assertRaises(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Optional[int]	=	Accelerator(cpu=_SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	:					Any	=	Accelerator()
    UpperCAmelCase	:					List[str]	=	GradientState()
    assert state.num_steps == 1
    UpperCAmelCase	:					Optional[Any]	=	4
    assert state.num_steps == 4

    assert state.sync_gradients is True
    UpperCAmelCase	:					List[str]	=	False
    assert state.sync_gradients is False
    GradientState._reset_state()

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





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					List[str]	=	create_components()

    (
        (
        UpperCAmelCase
    )	,		(
        UpperCAmelCase
    )	,		(
        UpperCAmelCase
    )	,		(
        UpperCAmelCase
    )	,		(
        UpperCAmelCase
    )	,		
    )	:					Union[str, Any]	=	accelerator.prepare(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

    self.assertTrue(prepared_model in accelerator._models						)
    self.assertTrue(prepared_optimizer in accelerator._optimizers						)
    self.assertTrue(prepared_scheduler in accelerator._schedulers						)
    self.assertTrue(prepared_train_dl in accelerator._dataloaders						)
    self.assertTrue(prepared_valid_dl in accelerator._dataloaders						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					int	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Dict	=	create_components()
    accelerator.prepare(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    accelerator.free_memory()

    self.assertTrue(len(accelerator._models						) == 0						)
    self.assertTrue(len(accelerator._optimizers						) == 0						)
    self.assertTrue(len(accelerator._schedulers						) == 0						)
    self.assertTrue(len(accelerator._dataloaders						) == 0						)

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





    '''simple docstring'''




    PartialState._reset_state()

    # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
    def noop(*_SCREAMING_SNAKE_CASE      ,     **_SCREAMING_SNAKE_CASE						):
     pass

    with patch("""torch.cuda.set_device"""      ,     _SCREAMING_SNAKE_CASE						), patch_environment(ACCELERATE_TORCH_DEVICE="""cuda:64"""						):
     UpperCAmelCase	:					List[str]	=	Accelerator()
     self.assertEqual(str(accelerator.state.device						)      ,     """cuda:64"""						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      str:





    '''simple docstring'''




    UpperCAmelCase	:					Any	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					int	=	create_components()
    accelerator.prepare(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Any	=	get_signature(_SCREAMING_SNAKE_CASE						)

    with tempfile.TemporaryDirectory() as tmpdirname:
     accelerator.save_state(_SCREAMING_SNAKE_CASE						)

     # make sure random weights don't match
     load_random_weights(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) > 1E-3						)

     # make sure loaded weights match
     accelerator.load_state(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) < 1E-3						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					int	=	create_components()
    accelerator.prepare(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Tuple	=	get_signature(_SCREAMING_SNAKE_CASE						)

    # saving hook
    def save_config(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					int	=	{"""class_name""": models[0].__class__.__name__}

     with open(os.path.join(_SCREAMING_SNAKE_CASE      ,     """data.json"""						)      ,     """w"""						) as f:
      json.dump(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

        # loading hook
    def load_config(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						):
     with open(os.path.join(_SCREAMING_SNAKE_CASE      ,     """data.json"""						)      ,     """r"""						) as f:
      UpperCAmelCase	:					Union[str, Any]	=	json.load(_SCREAMING_SNAKE_CASE						)

     UpperCAmelCase	:					Union[str, Any]	=	config["""class_name"""]

    UpperCAmelCase	:					Union[str, Any]	=	accelerator.register_save_state_pre_hook(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					List[str]	=	accelerator.register_load_state_pre_hook(_SCREAMING_SNAKE_CASE						)

    with tempfile.TemporaryDirectory() as tmpdirname:
     accelerator.save_state(_SCREAMING_SNAKE_CASE						)

     # make sure random weights don't match with hooks
     load_random_weights(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) > 1E-3						)

     # random class name to verify correct one is loaded
     UpperCAmelCase	:					Union[str, Any]	=	"""random"""

     # make sure loaded weights match with hooks
     accelerator.load_state(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) < 1E-3						)

     # mode.class_name is loaded from config
     self.assertTrue(model.class_name == model.__class__.__name__						)

    # remove hooks
    save_hook.remove()
    load_hook.remove()

    with tempfile.TemporaryDirectory() as tmpdirname:
     accelerator.save_state(_SCREAMING_SNAKE_CASE						)

     # make sure random weights don't match with hooks removed
     load_random_weights(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) > 1E-3						)

     # random class name to verify correct one is loaded
     UpperCAmelCase	:					List[Any]	=	"""random"""

     # make sure loaded weights match with hooks removed
     accelerator.load_state(_SCREAMING_SNAKE_CASE						)
     self.assertTrue(abs(model_signature - get_signature(_SCREAMING_SNAKE_CASE						)						) < 1E-3						)

     # mode.class_name is NOT loaded from config
     self.assertTrue(model.class_name != model.__class__.__name__						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      str:





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Optional[int]	=	create_components()
    UpperCAmelCase	:					List[str]	=	None

    # This should work
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Optional[Any]	=	accelerator.prepare(
        _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    self.assertTrue(dummy_obj is None						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					List[str]	=	Accelerator()
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					List[Any]	=	create_components()
    UpperCAmelCase	:					Union[str, Any]	=	[1, 2, 3]

    # This should work
    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Optional[Any]	=	accelerator.prepare(
        _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Dummy object should have `_is_accelerate_prepared` set to `True`"""      ,     )
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Model is missing `_is_accelerator_prepared` or is set to `False`"""      ,     )
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Optimizer is missing `_is_accelerator_prepared` or is set to `False`"""      ,     )
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Scheduler is missing `_is_accelerator_prepared` or is set to `False`"""      ,     )
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"""      ,     )
    self.assertEqual(
        getattr(_SCREAMING_SNAKE_CASE      ,     """_is_accelerate_prepared"""      ,     _SCREAMING_SNAKE_CASE						)      ,     _SCREAMING_SNAKE_CASE      ,     """Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"""      ,     )

   @slow
   @require_bnb
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    from transformers import AutoModelForCausalLM

    UpperCAmelCase	:					Any	=	AutoModelForCausalLM.from_pretrained(
        """EleutherAI/gpt-neo-125m"""      ,     load_in_abit=_SCREAMING_SNAKE_CASE      ,     device_map={"""""": 0}      ,     )
    UpperCAmelCase	:					Union[str, Any]	=	Accelerator()

    # This should work
    UpperCAmelCase	:					str	=	accelerator.prepare(_SCREAMING_SNAKE_CASE						)

   @slow
   @require_bnb
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    from transformers import AutoModelForCausalLM

    UpperCAmelCase	:					Any	=	Accelerator()

    with init_empty_weights():
     UpperCAmelCase	:					Optional[int]	=	AutoModelForCausalLM.from_pretrained(
         """EleutherAI/gpt-neo-125m"""      ,     )
     model.tie_weights()
     UpperCAmelCase	:					Any	=	infer_auto_device_map(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Optional[int]	=	"""cpu"""

    UpperCAmelCase	:					List[str]	=	AutoModelForCausalLM.from_pretrained(
        """EleutherAI/gpt-neo-125m"""      ,     device_map=_SCREAMING_SNAKE_CASE      ,     load_in_abit=_SCREAMING_SNAKE_CASE      ,     llm_inta_enable_fpaa_cpu_offload=_SCREAMING_SNAKE_CASE						)

    # This should not work and get value error
    with self.assertRaises(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Union[str, Any]	=	accelerator.prepare(_SCREAMING_SNAKE_CASE						)

   @slow
   @require_bnb
   @require_multi_gpu
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    from transformers import AutoModelForCausalLM

    UpperCAmelCase	:					Dict	=	{"""distributed_type""": DistributedType.MULTI_GPU}

    with init_empty_weights():
     UpperCAmelCase	:					Tuple	=	AutoModelForCausalLM.from_pretrained(
         """EleutherAI/gpt-neo-125m"""      ,     )
     model.tie_weights()
     UpperCAmelCase	:					List[str]	=	infer_auto_device_map(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					List[Any]	=	1

    UpperCAmelCase	:					Optional[int]	=	AutoModelForCausalLM.from_pretrained(
        """EleutherAI/gpt-neo-125m"""      ,     load_in_abit=_SCREAMING_SNAKE_CASE      ,     device_map=_SCREAMING_SNAKE_CASE      ,     )
    UpperCAmelCase	:					Optional[int]	=	Accelerator()

    # This should not work and get value error
    with self.assertRaises(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					Union[str, Any]	=	accelerator.prepare(_SCREAMING_SNAKE_CASE						)

    PartialState._reset_state()

   @slow
   @require_bnb
   @require_multi_gpu
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    from transformers import AutoModelForCausalLM

    with init_empty_weights():
     UpperCAmelCase	:					List[str]	=	AutoModelForCausalLM.from_pretrained(
         """EleutherAI/gpt-neo-125m"""      ,     )
     UpperCAmelCase	:					str	=	infer_auto_device_map(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					int	=	1

    UpperCAmelCase	:					int	=	AutoModelForCausalLM.from_pretrained(
        """EleutherAI/gpt-neo-125m"""      ,     load_in_abit=_SCREAMING_SNAKE_CASE      ,     device_map=_SCREAMING_SNAKE_CASE      ,     )
    UpperCAmelCase	:					int	=	Accelerator()

    # This should work
    UpperCAmelCase	:					List[str]	=	accelerator.prepare(_SCREAMING_SNAKE_CASE						)




   @require_cuda
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	torch.nn.Linear(10      ,     10						)
    UpperCAmelCase	:					int	=	torch.optim.SGD(model.parameters()      ,     lr=0.01						)
    UpperCAmelCase	:					str	=	Accelerator(cpu=_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					int	=	accelerator.prepare(_SCREAMING_SNAKE_CASE						)








import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
 import faiss





@require_faiss
class 			__a  (    A__       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  tempfile.mkdtemp()
      UpperCamelCase__			:   Any       				=  8

      # DPR tok
      UpperCamelCase__			:   List[Any]       				=  [
          "[UNK]",
          "[CLS]",
          "[SEP]",
          "[PAD]",
          "[MASK]",
          "want",
          "##want",
          "##ed",
          "wa",
          "un",
          "runn",
          "##ing",
          ",",
          "low",
          "lowest",
      ]
      UpperCamelCase__			:   Any       				=  os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)
      os.makedirs(SCREAMING_SNAKE_CASE							,   exist_ok=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[int]       				=  os.path.join(SCREAMING_SNAKE_CASE							,   DPR_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]						)						)

      # BART tok
      UpperCamelCase__			:   str       				=  [
          "l",
          "o",
          "w",
          "e",
          "r",
          "s",
          "t",
          "i",
          "d",
          "n",
          "\u0120",
          "\u0120l",
          "\u0120n",
          "\u0120lo",
          "\u0120low",
          "er",
          "\u0120lowest",
          "\u0120newer",
          "\u0120wider",
          "<unk>",
      ]
      UpperCamelCase__			:   int       				=  dict(zip(SCREAMING_SNAKE_CASE							,   range(len(SCREAMING_SNAKE_CASE						)						)						)						)
      UpperCamelCase__			:   Optional[Any]       				=  ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
      UpperCamelCase__			:   Union[str, Any]       				=  {"unk_token": "<unk>"}

      UpperCamelCase__			:   Dict       				=  os.path.join(self.tmpdirname							,   "bart_tokenizer"						)
      os.makedirs(SCREAMING_SNAKE_CASE							,   exist_ok=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   List[str]       				=  os.path.join(SCREAMING_SNAKE_CASE							,   BART_VOCAB_FILES_NAMES["vocab_file"]						)
      UpperCamelCase__			:   str       				=  os.path.join(SCREAMING_SNAKE_CASE							,   BART_VOCAB_FILES_NAMES["merges_file"]						)
      with open(self.vocab_file							,   "w"							,   encoding="utf-8"						) as fp:
           fp.write(json.dumps(SCREAMING_SNAKE_CASE						) + "\n"						)
      with open(self.merges_file							,   "w"							,   encoding="utf-8"						) as fp:
           fp.write("\n".join(SCREAMING_SNAKE_CASE						)						)







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)						)







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "bart_tokenizer"						)						)







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      shutil.rmtree(self.tmpdirname						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  Dataset.from_dict(
          {
              "id": ["0", "1"],
              "text": ["foo", "bar"],
              "title": ["Foo", "Bar"],
              "embeddings": [np.ones(self.retrieval_vector_size						), 2 * np.ones(self.retrieval_vector_size						)],
          }						)
      dataset.add_faiss_index("embeddings"							,   string_factory="Flat"							,   metric_type=faiss.METRIC_INNER_PRODUCT						)
      return dataset







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  self.get_dummy_dataset()
      UpperCamelCase__			:   List[Any]       				=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size							,   question_encoder=DPRConfig().to_dict()							,   generator=BartConfig().to_dict()							,   )
      with patch("transformers.models.rag.retrieval_rag.load_dataset"						) as mock_load_dataset:
           UpperCamelCase__			:   str       				=  dataset
           UpperCamelCase__			:   Optional[int]       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   )
      return retriever







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   bool						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.get_dummy_dataset()
      UpperCamelCase__			:   Tuple       				=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size							,   question_encoder=DPRConfig().to_dict()							,   generator=BartConfig().to_dict()							,   index_name="custom"							,   )
      if from_disk:
           UpperCamelCase__			:   Optional[int]       				=  os.path.join(self.tmpdirname							,   "dataset"						)
           UpperCamelCase__			:   List[str]       				=  os.path.join(self.tmpdirname							,   "index.faiss"						)
           dataset.get_index("embeddings"						).save(os.path.join(self.tmpdirname							,   "index.faiss"						)						)
           dataset.drop_index("embeddings"						)
           dataset.save_to_disk(os.path.join(self.tmpdirname							,   "dataset"						)						)
           del dataset
           UpperCamelCase__			:   str       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   )
      else:
           UpperCamelCase__			:   List[Any]       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   index=CustomHFIndex(config.retrieval_vector_size							,   SCREAMING_SNAKE_CASE						)							,   )
      return retriever







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  Dataset.from_dict(
          {
              "id": ["0", "1"],
              "text": ["foo", "bar"],
              "title": ["Foo", "Bar"],
              "embeddings": [np.ones(self.retrieval_vector_size + 1						), 2 * np.ones(self.retrieval_vector_size + 1						)],
          }						)
      dataset.add_faiss_index("embeddings"							,   string_factory="Flat"							,   metric_type=faiss.METRIC_INNER_PRODUCT						)

      UpperCamelCase__			:   List[str]       				=  os.path.join(self.tmpdirname							,   "hf_bert_base.hnswSQ8_correct_phi_128.c_index"						)
      dataset.save_faiss_index("embeddings"							,   index_file_name + ".index.dpr"						)
      pickle.dump(dataset["id"]							,   open(index_file_name + ".index_meta.dpr"							,   "wb"						)						)

      UpperCamelCase__			:   Optional[int]       				=  os.path.join(self.tmpdirname							,   "psgs_w100.tsv.pkl"						)
      UpperCamelCase__			:   Tuple       				=  {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
      pickle.dump(SCREAMING_SNAKE_CASE							,   open(SCREAMING_SNAKE_CASE							,   "wb"						)						)

      UpperCamelCase__			:   List[str]       				=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size							,   question_encoder=DPRConfig().to_dict()							,   generator=BartConfig().to_dict()							,   index_name="legacy"							,   index_path=self.tmpdirname							,   )
      UpperCamelCase__			:   List[str]       				=  RagRetriever(
          SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()						)
      return retriever







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  1
      UpperCamelCase__			:   Tuple       				=  self.get_dummy_canonical_hf_index_retriever()
      UpperCamelCase__			:   str       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      self.assertEqual(doc_dicts[0]["id"][0]							,   "1"						)  # max inner product is reached with second doc
      self.assertEqual(doc_dicts[1]["id"][0]							,   "0"						)  # max inner product is reached with first doc
      self.assertListEqual(doc_ids.tolist()							,   [[1], [0]]						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  self.get_dummy_canonical_hf_index_retriever()
      with tempfile.TemporaryDirectory() as tmp_dirname:
           with patch("transformers.models.rag.retrieval_rag.load_dataset"						) as mock_load_dataset:
                UpperCamelCase__			:   Optional[int]       				=  self.get_dummy_dataset()
                retriever.save_pretrained(SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   Any       				=  RagRetriever.from_pretrained(SCREAMING_SNAKE_CASE						)
                self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   str       				=  np.array(
                    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
                UpperCamelCase__			:   Any       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=1						)
                self.assertTrue(out is not None						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  1
      UpperCamelCase__			:   Dict       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      self.assertEqual(doc_dicts[0]["id"][0]							,   "1"						)  # max inner product is reached with second doc
      self.assertEqual(doc_dicts[1]["id"][0]							,   "0"						)  # max inner product is reached with first doc
      self.assertListEqual(doc_ids.tolist()							,   [[1], [0]]						)







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   Tuple       				=  RagRetriever.from_pretrained(SCREAMING_SNAKE_CASE						)
           self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   Optional[int]       				=  np.array(
               [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
           UpperCamelCase__			:   Any       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=1						)
           self.assertTrue(out is not None						)







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  1
      UpperCamelCase__			:   List[str]       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Dict       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      self.assertEqual(doc_dicts[0]["id"][0]							,   "1"						)  # max inner product is reached with second doc
      self.assertEqual(doc_dicts[1]["id"][0]							,   "0"						)  # max inner product is reached with first doc
      self.assertListEqual(doc_ids.tolist()							,   [[1], [0]]						)







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   Optional[int]       				=  RagRetriever.from_pretrained(SCREAMING_SNAKE_CASE						)
           self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   Dict       				=  np.array(
               [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
           UpperCamelCase__			:   Union[str, Any]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=1						)
           self.assertTrue(out is not None						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  1
      UpperCamelCase__			:   Any       				=  self.get_dummy_legacy_index_retriever()
      UpperCamelCase__			:   Optional[int]       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   Tuple       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["text"]						)							,   SCREAMING_SNAKE_CASE						)
      self.assertEqual(doc_dicts[0]["text"][0]							,   "bar"						)  # max inner product is reached with second doc
      self.assertEqual(doc_dicts[1]["text"][0]							,   "foo"						)  # max inner product is reached with first doc
      self.assertListEqual(doc_ids.tolist()							,   [[1], [0]]						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.get_dummy_legacy_index_retriever()
      with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   str       				=  RagRetriever.from_pretrained(SCREAMING_SNAKE_CASE						)
           self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
           UpperCamelCase__			:   int       				=  np.array(
               [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
           UpperCamelCase__			:   int       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=1						)
           self.assertTrue(out is not None						)







 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      import torch

      UpperCamelCase__			:   Optional[Any]       				=  1
      UpperCamelCase__			:   Optional[int]       				=  self.get_dummy_canonical_hf_index_retriever()
      UpperCamelCase__			:   Optional[Any]       				=  [[5, 7], [10, 11]]
      UpperCamelCase__			:   List[str]       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__			:   int       				=  retriever(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   prefix=retriever.config.generator.prefix							,   n_docs=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  (
          out["context_input_ids"],
          out["context_attention_mask"],
          out["retrieved_doc_embeds"],
      )
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   np.ndarray						)

      UpperCamelCase__			:   List[Any]       				=  retriever(
          SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   prefix=retriever.config.generator.prefix							,   n_docs=SCREAMING_SNAKE_CASE							,   return_tensors="pt"							,   )
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[Any]       				=  (  # noqa: F841
          out["context_input_ids"],
          out["context_attention_mask"],
          out["retrieved_doc_embeds"],
          out["doc_ids"],
      )
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   torch.Tensor						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   torch.Tensor						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   torch.Tensor						)







 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  self.get_dpr_ctx_encoder_tokenizer()
      UpperCamelCase__			:   Union[str, Any]       				=  1
      UpperCamelCase__			:   Tuple       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      retriever.set_ctx_encoder_tokenizer(SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Tuple       				=  [[5, 7], [10, 11]]
      UpperCamelCase__			:   List[str]       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__			:   Optional[int]       				=  retriever(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   prefix=retriever.config.generator.prefix							,   n_docs=SCREAMING_SNAKE_CASE						)

      self.assertEqual(
          len(SCREAMING_SNAKE_CASE						)							,   6						)  # check whether the retriever output consist of 6 attributes including tokenized docs
      self.assertEqual(
          all(k in out for k in ("tokenized_doc_ids", "tokenized_doc_attention_mask")						)							,   SCREAMING_SNAKE_CASE						)  # check for doc token related keys in dictionary.