



'''simple docstring'''





import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from tax import checkpoints

from transformers import MTaConfig, UMTaEncoderModel, UMTaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()



def   UpperCAmelCase			(  a_	,						a_	,						a_						)					->    Tuple:



   """simple docstring"""






   return params[F"{prefix}/{prefix}/relpos_bias/rel_embedding"][:, i, :]



def   UpperCAmelCase			(  a_	,						a_	,						a_	,						a_="attention"						)					->    Optional[Any]:



   """simple docstring"""






   A_   :						Optional[Any]		= np.ascontiguousarray(params[F"{prefix}/{prefix}/{layer_name}/key/kernel"][:, i, :, :]						)
   A_   :						Any		= k_tmp.reshape(k_tmp.shape[0]	,						k_tmp.shape[1] * k_tmp.shape[2]						)
   A_   :						Optional[Any]		= np.ascontiguousarray(params[F"{prefix}/{prefix}/{layer_name}/out/kernel"][:, i, :, :]						)
   A_   :						str		= o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1]	,						o_tmp.shape[2]						)
   A_   :						Optional[Any]		= np.ascontiguousarray(params[F"{prefix}/{prefix}/{layer_name}/query/kernel"][:, i, :, :]						)
   A_   :						Tuple		= q_tmp.reshape(q_tmp.shape[0]	,						q_tmp.shape[1] * q_tmp.shape[2]						)
   A_   :						Dict		= np.ascontiguousarray(params[F"{prefix}/{prefix}/{layer_name}/value/kernel"][:, i, :, :]						)
   A_   :						List[str]		= v_tmp.reshape(v_tmp.shape[0]	,						v_tmp.shape[1] * v_tmp.shape[2]						)
   return k, o, q, v



def   UpperCAmelCase			(  a_	,						a_	,						a_	,						a_=False						)					->    Optional[Any]:



   """simple docstring"""






   if split_mlp_wi:
      A_   :						Optional[int]		= params[F"{prefix}/{prefix}/mlp/wi_0/kernel"][:, i, :]
      A_   :						List[Any]		= params[F"{prefix}/{prefix}/mlp/wi_1/kernel"][:, i, :]
      A_   :						Optional[Any]		= (wi_a, wi_a)
   else:
      A_   :						Any		= params[F"{prefix}/{prefix}/mlp/wi/kernel"][:, i, :]

   A_   :						str		= params[F"{prefix}/{prefix}/mlp/wo/kernel"][:, i, :]
   return wi, wo



def   UpperCAmelCase			(  a_	,						a_	,						a_	,						a_						)					->    str:



   """simple docstring"""






   return params[F"{prefix}/{prefix}/{layer_name}/scale"][:, i]



def   UpperCAmelCase			(  a_	,						*, a_	,						a_	,						a_ = False						)					->    Union[str, Any]:



   """simple docstring"""






   A_   :						Dict		= traverse_util.flatten_dict(variables["""target"""]						)
   A_   :						int		= {"""/""".join(__lowercase						): v for k, v in old.items()}

   # v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
   A_   :						List[str]		= """encoder/encoder/mlp/wi_0/kernel""" in old
   print("""Split MLP:"""	,						__lowercase						)

   A_   :						int		= collections.OrderedDict()

   # Shared embeddings.
   A_   :						Any		= old["""token_embedder/embedding"""]

   # Encoder.
   for i in range(__lowercase						):
      # Block i, layer 0 (Self Attention).
      A_   :						List[Any]		= tax_layer_norm_lookup(__lowercase	,						__lowercase	,						"""encoder"""	,						"""pre_attention_layer_norm"""						)
      A_		,   A_		,   A_		,   A_   :						List[str]		= tax_attention_lookup(__lowercase	,						__lowercase	,						"""encoder"""	,						"""attention"""						)
      A_   :						str		= layer_norm
      A_   :						int		= k.T
      A_   :						Dict		= o.T
      A_   :						Dict		= q.T
      A_   :						List[Any]		= v.T

      # Block i, layer 1 (MLP).
      A_   :						Union[str, Any]		= tax_layer_norm_lookup(__lowercase	,						__lowercase	,						"""encoder"""	,						"""pre_mlp_layer_norm"""						)
      A_		,   A_   :						Optional[Any]		= tax_mlp_lookup(__lowercase	,						__lowercase	,						"""encoder"""	,						__lowercase						)
      A_   :						Dict		= layer_norm
      if split_mlp_wi:
         A_   :						Union[str, Any]		= wi[0].T
         A_   :						Optional[int]		= wi[1].T
      else:
         A_   :						int		= wi.T
      A_   :						Any		= wo.T
      if scalable_attention:
         # convert the rel_embedding of each layer
         A_   :						Optional[int]		= tax_relpos_bias_lookup(
             __lowercase	,						__lowercase	,						"""encoder"""						).T

   A_   :						str		= old["""encoder/encoder_norm/scale"""]

   if not scalable_attention:
      A_   :						List[str]		= tax_relpos_bias_lookup(
          __lowercase	,						0	,						"""encoder"""						).T
      A_   :						Tuple		= tax_relpos_bias_lookup(
          __lowercase	,						0	,						"""decoder"""						).T

   if not is_encoder_only:
      # Decoder.
      for i in range(__lowercase						):
         # Block i, layer 0 (Self Attention).
         A_   :						Optional[int]		= tax_layer_norm_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						"""pre_self_attention_layer_norm"""						)
         A_		,   A_		,   A_		,   A_   :						List[str]		= tax_attention_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						"""self_attention"""						)
         A_   :						Tuple		= layer_norm
         A_   :						int		= k.T
         A_   :						str		= o.T
         A_   :						List[str]		= q.T
         A_   :						Optional[int]		= v.T

         # Block i, layer 1 (Cross Attention).
         A_   :						Union[str, Any]		= tax_layer_norm_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						"""pre_cross_attention_layer_norm"""						)
         A_		,   A_		,   A_		,   A_   :						Dict		= tax_attention_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						"""encoder_decoder_attention"""						)
         A_   :						List[Any]		= layer_norm
         A_   :						Optional[int]		= k.T
         A_   :						str		= o.T
         A_   :						Optional[int]		= q.T
         A_   :						Dict		= v.T

         # Block i, layer 2 (MLP).
         A_   :						Any		= tax_layer_norm_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						"""pre_mlp_layer_norm"""						)
         A_		,   A_   :						int		= tax_mlp_lookup(__lowercase	,						__lowercase	,						"""decoder"""	,						__lowercase						)
         A_   :						List[Any]		= layer_norm
         if split_mlp_wi:
            A_   :						int		= wi[0].T
            A_   :						List[str]		= wi[1].T
         else:
            A_   :						Optional[int]		= wi.T
         A_   :						Dict		= wo.T

         if scalable_attention:
            # convert the rel_embedding of each layer
            A_   :						str		= tax_relpos_bias_lookup(__lowercase	,						__lowercase	,						"""decoder"""						).T

      A_   :						Union[str, Any]		= old["""decoder/decoder_norm/scale"""]

      # LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
      if "decoder/logits_dense/kernel" in old:
         A_   :						Tuple		= old["""decoder/logits_dense/kernel"""].T

   return new



def   UpperCAmelCase			(  a_	,						a_						)					->    int:



   """simple docstring"""






   A_   :						str		= collections.OrderedDict([(k, torch.from_numpy(v.copy()						)) for (k, v) in converted_params.items()]						)

   # Add what is missing.
   if "encoder.embed_tokens.weight" not in state_dict:
      A_   :						Optional[int]		= state_dict["""shared.weight"""]

   if not is_encoder_only:
      if "decoder.embed_tokens.weight" not in state_dict:
         A_   :						Union[str, Any]		= state_dict["""shared.weight"""]

      if "lm_head.weight" not in state_dict:  # For old 1.0 models.
         print("""Using shared word embeddings as lm_head."""						)
         A_   :						List[str]		= state_dict["""shared.weight"""]

   return state_dict



def   UpperCAmelCase			(  a_	,						a_	,						a_	,						a_	,						a_						)					->    Union[str, Any]:



   """simple docstring"""






   A_   :						Any		= checkpoints.load_tax_checkpoint(__lowercase						)
   A_   :						Dict		= convert_tax_to_pytorch(
       __lowercase	,						num_layers=config.num_layers	,						is_encoder_only=__lowercase	,						scalable_attention=__lowercase						)
   A_   :						Any		= make_state_dict(__lowercase	,						__lowercase						)
   model.load_state_dict(__lowercase	,						strict=__lowercase						)



def   UpperCAmelCase			(  a_	,						a_	,						a_	,						a_ = False	,						a_ = False	,						)					->    Optional[int]:



   """simple docstring"""






   A_   :						Tuple		= MTaConfig.from_json_file(__lowercase						)
   print(F"Building PyTorch model from configuration: {config}"						)
   # Non-v1.1 checkpoints could also use T5Model, but this works for all.
   # The v1.0 checkpoints will simply have an LM head that is the word embeddings.
   if is_encoder_only:
      A_   :						str		= UMTaEncoderModel(__lowercase						)
   else:
      A_   :						List[str]		= UMTaForConditionalGeneration(__lowercase						)

   # Load weights from tf checkpoint
   load_tax_weights_in_ta(__lowercase	,						__lowercase	,						__lowercase	,						__lowercase	,						__lowercase						)

   # Save pytorch-model
   print(F"Save PyTorch model to {pytorch_dump_path}"						)
   model.save_pretrained(__lowercase						)

   # Verify that we can load the checkpoint.
   model.from_pretrained(__lowercase						)
   print("""Done"""						)


if __name__ == "__main__":
  UpperCamelCase__  :			List[Any]          =					argparse.ArgumentParser(description='Converts a native T5X checkpoint into a PyTorch checkpoint.')
  # Required parameters
  parser.add_argument(
      '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path to the T5X checkpoint.'
  )
  parser.add_argument(
      '--config_file',
      default=None,
      type=str,
      required=True,
      help='The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.',
  )
  parser.add_argument(
      '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
  )
  parser.add_argument(
      '--is_encoder_only', action='store_true', help='Check if the model is encoder-decoder model', default=False
  )
  parser.add_argument(
      '--scalable_attention',
      action='store_true',
      help='Whether the model uses scaled attention (umt5 model)',
      default=False,
  )
  UpperCamelCase__  :			List[str]          =					parser.parse_args()
  convert_tax_checkpoint_to_pytorch(
      args.tax_checkpoint_path,
      args.config_file,
      args.pytorch_dump_path,
      args.is_encoder_only,
      args.scalable_attention,
  )







'''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.
