



'''simple docstring'''





import argparse
import json
import logging
import os
import shutil
import sys
import tempfile
import unittest
from unittest import mock

import torch
from accelerate.utils import write_basic_config

from transformers.testing_utils import TestCasePlus, get_gpu_count, run_command, slow, torch_device
from transformers.utils import is_apex_available


logging.basicConfig(level=logging.DEBUG)

UpperCamelCase__  :			Any          =					logging.getLogger()



def   UpperCAmelCase			(  )					->    List[Any]:



   """simple docstring"""






   A_   :						Dict		= argparse.ArgumentParser()
   parser.add_argument("""-f"""						)
   A_   :						Tuple		= parser.parse_args()
   return args.f



def   UpperCAmelCase			(  a_						)					->    Any:



   """simple docstring"""






   A_   :						Optional[Any]		= {}
   A_   :						Optional[int]		= os.path.join(__lowercase	,						"""all_results.json"""						)
   if os.path.exists(__lowercase						):
      with open(__lowercase	,						"""r"""						) as f:
         A_   :						Optional[Any]		= json.load(__lowercase						)
   else:
      raise ValueError(F"can\'t find {path}"						)
   return results



def   UpperCAmelCase			(  )					->    Dict:



   """simple docstring"""






   A_   :						Optional[Any]		= torch.cuda.is_available() and torch_device == """cuda"""
   return is_using_cuda and is_apex_available()


UpperCamelCase__  :			str          =					logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 				_lowerCAmelCase     (   __lowerCamelCase				):


 """simple docstring"""

 @classmethod
 def 	UpperCAmelCase_    (	cls   )  ->  Any:
    # Write Accelerate config, will pick up on CPU, GPU, and multi-GPU
    A_   :						Tuple		= tempfile.mkdtemp()
    A_   :						Dict		= os.path.join(cls.tmpdir ,  """default_config.yml"""   )
    write_basic_config(save_location=cls.configPath   )
    A_   :						Dict		= ["""accelerate""", """launch""", """--config_file""", cls.configPath]







 @classmethod
 def 	UpperCAmelCase_    (	cls   )  ->  Optional[int]:
    shutil.rmtree(cls.tmpdir   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Tuple:
    A_   :						List[str]		= self.get_auto_remove_tmp_dir()
    A_   :						List[Any]		= F"\n            {self.examples_dir}/pytorch/text-classification/run_glue_no_trainer.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --seed=42\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    if is_cuda_and_apex_available():
       testargs.append("""--fp16"""   )

    run_command(self._launch_args + testargs   )
    A_   :						Dict		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_accuracy"""] ,  0.75   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """glue_no_trainer"""   )   )   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Optional[Any]:
    A_   :						Union[str, Any]		= self.get_auto_remove_tmp_dir()
    A_   :						Union[str, Any]		= F"\n            {self.examples_dir}/pytorch/language-modeling/run_clm_no_trainer.py\n            --model_name_or_path distilgpt2\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --block_size 128\n            --per_device_train_batch_size 5\n            --per_device_eval_batch_size 5\n            --num_train_epochs 2\n            --output_dir {tmp_dir}\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    if torch.cuda.device_count() > 1:
       # Skipping because there are not enough batches to train the model + would need a drop_last to work.
       return

    run_command(self._launch_args + testargs   )
    A_   :						Tuple		= get_results(__A   )
    self.assertLess(result["""perplexity"""] ,  100   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """clm_no_trainer"""   )   )   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Union[str, Any]:
    A_   :						Optional[Any]		= self.get_auto_remove_tmp_dir()
    A_   :						Dict		= F"\n            {self.examples_dir}/pytorch/language-modeling/run_mlm_no_trainer.py\n            --model_name_or_path distilroberta-base\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --output_dir {tmp_dir}\n            --num_train_epochs=1\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Optional[int]		= get_results(__A   )
    self.assertLess(result["""perplexity"""] ,  42   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """mlm_no_trainer"""   )   )   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  str:
    # with so little data distributed training needs more epochs to get the score on par with 0/1 gpu
    A_   :						Dict		= 7 if get_gpu_count() > 1 else 2

    A_   :						Tuple		= self.get_auto_remove_tmp_dir()
    A_   :						List[str]		= F"\n            {self.examples_dir}/pytorch/token-classification/run_ner_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --train_file tests/fixtures/tests_samples/conll/sample.json\n            --validation_file tests/fixtures/tests_samples/conll/sample.json\n            --output_dir {tmp_dir}\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=2\n            --num_train_epochs={epochs}\n            --seed 7\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Dict		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_accuracy"""] ,  0.75   )
    self.assertLess(result["""train_loss"""] ,  0.5   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """ner_no_trainer"""   )   )   )







 @unittest.skip(reason="""Fix me @muellerzr"""   )
 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Tuple:
    A_   :						int		= self.get_auto_remove_tmp_dir()
    A_   :						int		= F"\n            {self.examples_dir}/pytorch/question-answering/run_qa_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --version_2_with_negative\n            --train_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --output_dir {tmp_dir}\n            --seed=42\n            --max_train_steps=10\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Union[str, Any]		= get_results(__A   )
    # Because we use --version_2_with_negative the testing script uses SQuAD v2 metrics.
    self.assertGreaterEqual(result["""eval_f1"""] ,  28   )
    self.assertGreaterEqual(result["""eval_exact"""] ,  28   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """qa_no_trainer"""   )   )   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Union[str, Any]:
    A_   :						Tuple		= self.get_auto_remove_tmp_dir()
    A_   :						Optional[Any]		= F"\n            {self.examples_dir}/pytorch/multiple-choice/run_swag_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --train_file tests/fixtures/tests_samples/swag/sample.json\n            --validation_file tests/fixtures/tests_samples/swag/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=20\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Optional[Any]		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_accuracy"""] ,  0.8   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """swag_no_trainer"""   )   )   )







 @slow
 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  List[Any]:
    A_   :						List[str]		= self.get_auto_remove_tmp_dir()
    A_   :						Dict		= F"\n            {self.examples_dir}/pytorch/summarization/run_summarization_no_trainer.py\n            --model_name_or_path t5-small\n            --train_file tests/fixtures/tests_samples/xsum/sample.json\n            --validation_file tests/fixtures/tests_samples/xsum/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=50\n            --num_warmup_steps=8\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Dict		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_rouge1"""] ,  10   )
    self.assertGreaterEqual(result["""eval_rouge2"""] ,  2   )
    self.assertGreaterEqual(result["""eval_rougeL"""] ,  7   )
    self.assertGreaterEqual(result["""eval_rougeLsum"""] ,  7   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """summarization_no_trainer"""   )   )   )







 @slow
 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  List[str]:
    A_   :						List[str]		= self.get_auto_remove_tmp_dir()
    A_   :						Optional[int]		= F"\n            {self.examples_dir}/pytorch/translation/run_translation_no_trainer.py\n            --model_name_or_path sshleifer/student_marian_en_ro_6_1\n            --source_lang en\n            --target_lang ro\n            --train_file tests/fixtures/tests_samples/wmt16/sample.json\n            --validation_file tests/fixtures/tests_samples/wmt16/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=50\n            --num_warmup_steps=8\n            --num_beams=6\n            --learning_rate=3e-3\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --source_lang en_XX\n            --target_lang ro_RO\n            --checkpointing_steps epoch\n            --with_tracking\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						Optional[int]		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_bleu"""] ,  30   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """epoch_0"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """translation_no_trainer"""   )   )   )







 @slow
 def 	UpperCAmelCase_    (	self   )  ->  str:
    A_   :						Any		= logging.StreamHandler(sys.stdout   )
    logger.addHandler(__A   )

    A_   :						Optional[int]		= self.get_auto_remove_tmp_dir()
    A_   :						Optional[int]		= F"\n            {self.examples_dir}/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py\n            --dataset_name huggingface/semantic-segmentation-test-sample\n            --output_dir {tmp_dir}\n            --max_train_steps=10\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n        ".split()

    run_command(self._launch_args + testargs   )
    A_   :						str		= get_results(__A   )
    self.assertGreaterEqual(result["""eval_overall_accuracy"""] ,  0.10   )







 @mock.patch.dict(os.environ ,  {"""WANDB_MODE""": """offline"""}   )
 def 	UpperCAmelCase_    (	self   )  ->  Any:
    A_   :						Optional[Any]		= self.get_auto_remove_tmp_dir()
    A_   :						Tuple		= F"\n            {self.examples_dir}/pytorch/image-classification/run_image_classification_no_trainer.py\n            --model_name_or_path google/vit-base-patch16-224-in21k\n            --dataset_name hf-internal-testing/cats_vs_dogs_sample\n            --learning_rate 1e-4\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 1\n            --max_train_steps 2\n            --train_val_split 0.1\n            --seed 42\n            --output_dir {tmp_dir}\n            --with_tracking\n            --checkpointing_steps 1\n        ".split()

    if is_cuda_and_apex_available():
       testargs.append("""--fp16"""   )

    run_command(self._launch_args + testargs   )
    A_   :						Tuple		= get_results(__A   )
    # The base model scores a 25%
    self.assertGreaterEqual(result["""eval_accuracy"""] ,  0.6   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """step_1"""   )   )   )
    self.assertTrue(os.path.exists(os.path.join(__A ,  """image_classification_no_trainer"""   )   )   )







'''simple docstring'''
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      snake_case     (    __lowerCamelCase       ):




 """simple docstring"""



 def        _lowerCamelCase					(       self       :							Any	):
        __UpperCamelCase       =  tempfile.mkdtemp()
        __UpperCamelCase       =  8

        # DPR tok
        __UpperCamelCase       =  [
            '[UNK]',
            '[CLS]',
            '[SEP]',
            '[PAD]',
            '[MASK]',
            'want',
            '##want',
            '##ed',
            'wa',
            'un',
            'runn',
            '##ing',
            ',',
            'low',
            'lowest',
        ]
        __UpperCamelCase       =  os.path.join(self.tmpdirname							,	'dpr_tokenizer'	)
        os.makedirs(__A							,	exist_ok=__A	)
        __UpperCamelCase       =  os.path.join(__A							,	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       =  [
            'l',
            'o',
            'w',
            'e',
            'r',
            's',
            't',
            'i',
            'd',
            'n',
            '\u0120',
            '\u0120l',
            '\u0120n',
            '\u0120lo',
            '\u0120low',
            'er',
            '\u0120lowest',
            '\u0120newer',
            '\u0120wider',
            '<unk>',
        ]
        __UpperCamelCase       =  dict(zip(__A							,	range(len(__A	)	)	)	)
        __UpperCamelCase       =  ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
        __UpperCamelCase       =  {'unk_token': '<unk>'}

        __UpperCamelCase       =  os.path.join(self.tmpdirname							,	'bart_tokenizer'	)
        os.makedirs(__A							,	exist_ok=__A	)
        __UpperCamelCase       =  os.path.join(__A							,	BART_VOCAB_FILES_NAMES['vocab_file']	)
        __UpperCamelCase       =  os.path.join(__A							,	BART_VOCAB_FILES_NAMES['merges_file']	)
        with open(self.vocab_file							,	'w'							,	encoding='utf-8'	) as fp:
               fp.write(json.dumps(__A	) + '\n'	)
        with open(self.merges_file							,	'w'							,	encoding='utf-8'	) as fp:
               fp.write('\n'.join(__A	)	)






 def        _lowerCamelCase					(       self       :							Tuple	):
        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	'dpr_tokenizer'	)	)






 def        _lowerCamelCase					(       self       :							Optional[int]	):
        return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	'dpr_tokenizer'	)	)






 def        _lowerCamelCase					(       self       :							Union[str, Any]	):
        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	'bart_tokenizer'	)	)






 def        _lowerCamelCase					(       self       :							str	):
        shutil.rmtree(self.tmpdirname	)






 def        _lowerCamelCase					(       self       :							Dict	):
        __UpperCamelCase       =  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        _lowerCamelCase					(       self       :							Tuple	):
        __UpperCamelCase       =  self.get_dummy_dataset()
        __UpperCamelCase       =  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       =  dataset
               __UpperCamelCase       =  RagRetriever(
                   __A							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	)
        return retriever






 def        _lowerCamelCase					(       self       :							Any							,	__A       :							bool	):
        __UpperCamelCase       =  self.get_dummy_dataset()
        __UpperCamelCase       =  RagConfig(
            retrieval_vector_size=self.retrieval_vector_size							,	question_encoder=DPRConfig().to_dict()							,	generator=BartConfig().to_dict()							,	index_name='custom'							,	)
        if from_disk:
               __UpperCamelCase       =  os.path.join(self.tmpdirname							,	'dataset'	)
               __UpperCamelCase       =  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       =  RagRetriever(
                   __A							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	)
        else:
               __UpperCamelCase       =  RagRetriever(
                   __A							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	index=CustomHFIndex(config.retrieval_vector_size							,	__A	)							,	)
        return retriever






 def        _lowerCamelCase					(       self       :							int	):
        __UpperCamelCase       =  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       =  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       =  os.path.join(self.tmpdirname							,	'psgs_w100.tsv.pkl'	)
        __UpperCamelCase       =  {sample['id']: [sample['text'], sample['title']] for sample in dataset}
        pickle.dump(__A							,	open(__A							,	'wb'	)	)

        __UpperCamelCase       =  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       =  RagRetriever(
            __A							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()	)
        return retriever






 def        _lowerCamelCase					(       self       :							List[str]	):
        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_canonical_hf_index_retriever()
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=__A	)
        self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size)	)
        self.assertEqual(len(__A	)							,	2	)
        self.assertEqual(sorted(doc_dicts[0]	)							,	['embeddings', 'id', 'text', 'title']	)
        self.assertEqual(len(doc_dicts[0]['id']	)							,	__A	)
        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        _lowerCamelCase					(       self       :							Optional[Any]	):
        __UpperCamelCase       =  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       =  self.get_dummy_dataset()
                      retriever.save_pretrained(__A	)
                      __UpperCamelCase       =  RagRetriever.from_pretrained(__A	)
                      self.assertIsInstance(__A							,	__A	)
                      __UpperCamelCase       =  np.array(
                          [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
                      __UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=1	)
                      self.assertTrue(out is not None	)






 def        _lowerCamelCase					(       self       :							Optional[int]	):
        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_custom_hf_index_retriever(from_disk=__A	)
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=__A	)
        self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size)	)
        self.assertEqual(len(__A	)							,	2	)
        self.assertEqual(sorted(doc_dicts[0]	)							,	['embeddings', 'id', 'text', 'title']	)
        self.assertEqual(len(doc_dicts[0]['id']	)							,	__A	)
        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        _lowerCamelCase					(       self       :							str	):
        __UpperCamelCase       =  self.get_dummy_custom_hf_index_retriever(from_disk=__A	)
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__A	)
               __UpperCamelCase       =  RagRetriever.from_pretrained(__A	)
               self.assertIsInstance(__A							,	__A	)
               __UpperCamelCase       =  np.array(
                   [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
               __UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=1	)
               self.assertTrue(out is not None	)






 def        _lowerCamelCase					(       self       :							Optional[Any]	):
        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_custom_hf_index_retriever(from_disk=__A	)
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=__A	)
        self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size)	)
        self.assertEqual(len(__A	)							,	2	)
        self.assertEqual(sorted(doc_dicts[0]	)							,	['embeddings', 'id', 'text', 'title']	)
        self.assertEqual(len(doc_dicts[0]['id']	)							,	__A	)
        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        _lowerCamelCase					(       self       :							Any	):
        __UpperCamelCase       =  self.get_dummy_custom_hf_index_retriever(from_disk=__A	)
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__A	)
               __UpperCamelCase       =  RagRetriever.from_pretrained(__A	)
               self.assertIsInstance(__A							,	__A	)
               __UpperCamelCase       =  np.array(
                   [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
               __UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=1	)
               self.assertTrue(out is not None	)






 def        _lowerCamelCase					(       self       :							Dict	):
        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_legacy_index_retriever()
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=__A	)
        self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size)	)
        self.assertEqual(len(__A	)							,	2	)
        self.assertEqual(sorted(doc_dicts[0]	)							,	['text', 'title']	)
        self.assertEqual(len(doc_dicts[0]['text']	)							,	__A	)
        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        _lowerCamelCase					(       self       :							str	):
        __UpperCamelCase       =  self.get_dummy_legacy_index_retriever()
        with tempfile.TemporaryDirectory() as tmp_dirname:
               retriever.save_pretrained(__A	)
               __UpperCamelCase       =  RagRetriever.from_pretrained(__A	)
               self.assertIsInstance(__A							,	__A	)
               __UpperCamelCase       =  np.array(
                   [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
               __UpperCamelCase       =  retriever.retrieve(__A							,	n_docs=1	)
               self.assertTrue(out is not None	)






 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        _lowerCamelCase					(       self       :							Optional[Any]	):
        import torch

        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_canonical_hf_index_retriever()
        __UpperCamelCase       =  [[5, 7], [1_0, 1_1]]
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       =  retriever(__A							,	__A							,	prefix=retriever.config.generator.prefix							,	n_docs=__A	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  (
            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(__A							,	__A	)
        self.assertIsInstance(__A							,	__A	)
        self.assertIsInstance(__A							,	np.ndarray	)

        __UpperCamelCase       =  retriever(
            __A							,	__A							,	prefix=retriever.config.generator.prefix							,	n_docs=__A							,	return_tensors='pt'							,	)
        __UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       ,						__UpperCamelCase       =  (  # 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(__A							,	torch.Tensor	)
        self.assertIsInstance(__A							,	torch.Tensor	)
        self.assertIsInstance(__A							,	torch.Tensor	)






 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def        _lowerCamelCase					(       self       :							Any	):
        __UpperCamelCase       =  self.get_dpr_ctx_encoder_tokenizer()
        __UpperCamelCase       =  1
        __UpperCamelCase       =  self.get_dummy_custom_hf_index_retriever(from_disk=__A	)
        retriever.set_ctx_encoder_tokenizer(__A	)

        __UpperCamelCase       =  [[5, 7], [1_0, 1_1]]
        __UpperCamelCase       =  np.array(
            [np.ones(self.retrieval_vector_size	), -np.ones(self.retrieval_vector_size	)]							,	dtype=np.floataa	)
        __UpperCamelCase       =  retriever(__A							,	__A							,	prefix=retriever.config.generator.prefix							,	n_docs=__A	)

        self.assertEqual(
            len(__A	)							,	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')	)							,	__A	)  # check for doc token related keys in dictionary.
