



"""simple docstring"""
import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


A:     str										=	logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm
      import torch_xla.debug.metrics as met



class    SCREAMING_SNAKE_CASE__						(      UpperCAmelCase__				):

   def __init__(     self   ,  *_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE=None   ,  **_SCREAMING_SNAKE_CASE		)						->   Tuple:





    '''simple docstring'''




    super().__init__(*_SCREAMING_SNAKE_CASE   ,  **_SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      Tuple			=			eval_examples
    UpperCAmelCase			:      List[Any]			=			post_process_function
    UpperCAmelCase			:      List[Any]			=			quant_trainer_args
    UpperCAmelCase			:      Dict			=			128  # default number of calibration samples

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE=None		)						->   Optional[int]:





    '''simple docstring'''




    if calib_dataset is None and self.calib_dataset is None:
     raise ValueError("""Trainer: calibration requires an calib_dataset."""		)
    UpperCAmelCase			:      List[str]			=			calib_dataset if calib_dataset is not None else self.calib_dataset

    UpperCAmelCase			:      Tuple			=			self._remove_unused_columns(_SCREAMING_SNAKE_CASE   ,  description="""Calibration"""		)

    return DataLoader(
        _SCREAMING_SNAKE_CASE   ,  batch_size=self.args.eval_batch_size   ,  collate_fn=self.data_collator   ,  drop_last=self.args.dataloader_drop_last   ,  num_workers=self.args.dataloader_num_workers   ,  pin_memory=self.args.dataloader_pin_memory   ,  shuffle=_SCREAMING_SNAKE_CASE   ,  )

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE=None		)						->   Tuple:





    '''simple docstring'''




    UpperCAmelCase			:      Union[str, Any]			=			self.train_dataset if calib_dataset is None else calib_dataset
    UpperCAmelCase			:      Tuple			=			self.get_calib_dataloader(_SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Optional[int]			=			self.model
    quant_trainer.configure_model(_SCREAMING_SNAKE_CASE   ,  self.quant_trainer_args   ,  calib=_SCREAMING_SNAKE_CASE		)
    model.eval()
    quant_trainer.enable_calibration(_SCREAMING_SNAKE_CASE		)

    logger.info("""***** Running calibration *****"""		)
    logger.info(F"  Num examples = {self.calib_num}"		)
    logger.info(F"  Batch size = {calib_dataloader.batch_size}"		)

    for step, inputs in enumerate(_SCREAMING_SNAKE_CASE		):
     # Prediction step
     UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase			:      List[Any]			=			self.prediction_step(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  prediction_loss_only=_SCREAMING_SNAKE_CASE		)
     if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
      break

    quant_trainer.finish_calibration(_SCREAMING_SNAKE_CASE   ,  self.quant_trainer_args		)
    UpperCAmelCase			:      Optional[int]			=			model

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE = "eval"		)						->   Union[str, Any]:





    '''simple docstring'''




    UpperCAmelCase			:      str			=			self.eval_dataset if eval_dataset is None else eval_dataset
    UpperCAmelCase			:      int			=			self.get_eval_dataloader(_SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      Dict			=			self.eval_examples if eval_examples is None else eval_examples

    # Temporarily disable metric computation, we will do it in the loop here.
    UpperCAmelCase			:      Dict			=			self.compute_metrics
    UpperCAmelCase			:      int			=			None
    UpperCAmelCase			:      Any			=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
    try:
     UpperCAmelCase			:      Dict			=			eval_loop(
         _SCREAMING_SNAKE_CASE   ,  description="""Evaluation"""   ,  prediction_loss_only=True if compute_metrics is None else None   ,  ignore_keys=_SCREAMING_SNAKE_CASE   ,  )
    finally:
     UpperCAmelCase			:      str			=			compute_metrics

    if self.post_process_function is not None and self.compute_metrics is not None:
     UpperCAmelCase			:      Optional[Any]			=			self.post_process_function(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  output.predictions		)
     UpperCAmelCase			:      str			=			self.compute_metrics(_SCREAMING_SNAKE_CASE		)

     # Prefix all keys with metric_key_prefix + '_'
     for key in list(metrics.keys()		):
      if not key.startswith(F"{metric_key_prefix}_"		):
       UpperCAmelCase			:      str			=			metrics.pop(_SCREAMING_SNAKE_CASE		)

     self.log(_SCREAMING_SNAKE_CASE		)
    else:
     UpperCAmelCase			:      Union[str, Any]			=			{}

    if self.args.tpu_metrics_debug or self.args.debug:
     # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
     xm.master_print(met.metrics_report()		)

    UpperCAmelCase			:      Tuple			=			self.callback_handler.on_evaluate(self.args   ,  self.state   ,  self.control   ,  _SCREAMING_SNAKE_CASE		)
    return metrics

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE = "test"		)						->   List[str]:





    '''simple docstring'''




    UpperCAmelCase			:      int			=			self.get_test_dataloader(_SCREAMING_SNAKE_CASE		)

    # Temporarily disable metric computation, we will do it in the loop here.
    UpperCAmelCase			:      str			=			self.compute_metrics
    UpperCAmelCase			:      int			=			None
    UpperCAmelCase			:      Union[str, Any]			=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
    try:
     UpperCAmelCase			:      Union[str, Any]			=			eval_loop(
         _SCREAMING_SNAKE_CASE   ,  description="""Prediction"""   ,  prediction_loss_only=True if compute_metrics is None else None   ,  ignore_keys=_SCREAMING_SNAKE_CASE   ,  )
    finally:
     UpperCAmelCase			:      Any			=			compute_metrics

    if self.post_process_function is None or self.compute_metrics is None:
     return output

    UpperCAmelCase			:      Optional[int]			=			self.post_process_function(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  output.predictions   ,  """predict"""		)
    UpperCAmelCase			:      str			=			self.compute_metrics(_SCREAMING_SNAKE_CASE		)

    # Prefix all keys with metric_key_prefix + '_'
    for key in list(metrics.keys()		):
     if not key.startswith(F"{metric_key_prefix}_"		):
      UpperCAmelCase			:      List[str]			=			metrics.pop(_SCREAMING_SNAKE_CASE		)

    return PredictionOutput(predictions=predictions.predictions   ,  label_ids=predictions.label_ids   ,  metrics=_SCREAMING_SNAKE_CASE		)




   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE="./"		)						->   List[str]:





    '''simple docstring'''




    UpperCAmelCase			:      Any			=			self.eval_dataset
    UpperCAmelCase			:      Optional[Any]			=			self.get_eval_dataloader(_SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Tuple			=			next(iter(_SCREAMING_SNAKE_CASE		)		)

    # saving device - to make it consistent
    UpperCAmelCase			:      List[Any]			=			torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""		)

    # convert to tuple
    UpperCAmelCase			:      List[Any]			=			tuple(v.to(_SCREAMING_SNAKE_CASE		) for k, v in batch.items()		)

    logger.info("""Converting model to be onnx compatible"""		)
    from pytorch_quantization.nn import TensorQuantizer

    UpperCAmelCase			:      List[Any]			=			True

    UpperCAmelCase			:      Any			=			self.model.to(_SCREAMING_SNAKE_CASE		)

    model.eval()
    model.float()

    UpperCAmelCase			:      str			=			model.module if hasattr(_SCREAMING_SNAKE_CASE   ,  """module"""		) else model
    quant_trainer.configure_model(_SCREAMING_SNAKE_CASE   ,  self.quant_trainer_args		)

    UpperCAmelCase			:      Tuple			=			os.path.join(_SCREAMING_SNAKE_CASE   ,  """model.onnx"""		)
    logger.info(F"exporting model to {output_model_file}"		)

    UpperCAmelCase			:      List[Any]			=			{0: """batch_size""", 1: """seq_len"""}

    torch.onnx.export(
        _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  export_params=_SCREAMING_SNAKE_CASE   ,  opset_version=13   ,  do_constant_folding=_SCREAMING_SNAKE_CASE   ,  input_names=["""input_ids""", """attention_mask""", """token_type_ids"""]   ,  output_names=["""output_start_logits""", """output_end_logits"""]   ,  dynamic_axes={
            """input_ids""": axes,
            """attention_mask""": axes,
            """token_type_ids""": axes,
            """output_start_logits""": axes,
            """output_end_logits""": axes,
        }   ,  verbose=_SCREAMING_SNAKE_CASE   ,  )
    logger.info("""onnx export finished"""		)





"""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    SCREAMING_SNAKE_CASE__						(      UpperCAmelCase__				):

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





    '''simple docstring'''




    UpperCAmelCase			:      str			=			tempfile.mkdtemp()
    UpperCAmelCase			:      str			=			8

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

    UpperCAmelCase			:      int			=			os.path.join(self.tmpdirname   ,  """bart_tokenizer"""		)
    os.makedirs(_SCREAMING_SNAKE_CASE   ,  exist_ok=_SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      Union[str, Any]			=			os.path.join(_SCREAMING_SNAKE_CASE   ,  BART_VOCAB_FILES_NAMES["""vocab_file"""]		)
    UpperCAmelCase			:      Union[str, Any]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   DPRQuestionEncoderTokenizer:





    '''simple docstring'''




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

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   DPRContextEncoderTokenizer:





    '''simple docstring'''




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

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   BartTokenizer:





    '''simple docstring'''




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

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





    '''simple docstring'''




    shutil.rmtree(self.tmpdirname		)

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





    '''simple docstring'''




    UpperCAmelCase			:      Dict			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase			:      Union[str, Any]			=			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			:      Tuple			=			dataset
     UpperCAmelCase			:      List[str]			=			RagRetriever(
         _SCREAMING_SNAKE_CASE   ,  question_encoder_tokenizer=self.get_dpr_tokenizer()   ,  generator_tokenizer=self.get_bart_tokenizer()   ,  )
    return retriever

   def 			SCREAMING_SNAKE_CASE  (     self   ,  _SCREAMING_SNAKE_CASE		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      List[Any]			=			self.get_dummy_dataset()
    UpperCAmelCase			:      Any			=			RagConfig(
        retrieval_vector_size=self.retrieval_vector_size   ,  question_encoder=DPRConfig().to_dict()   ,  generator=BartConfig().to_dict()   ,  index_name="""custom"""   ,  )
    if from_disk:
     UpperCAmelCase			:      str			=			os.path.join(self.tmpdirname   ,  """dataset"""		)
     UpperCAmelCase			:      Union[str, Any]			=			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			:      List[Any]			=			RagRetriever(
         _SCREAMING_SNAKE_CASE   ,  question_encoder_tokenizer=self.get_dpr_tokenizer()   ,  generator_tokenizer=self.get_bart_tokenizer()   ,  )
    else:
     UpperCAmelCase			:      Optional[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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





    '''simple docstring'''




    UpperCAmelCase			:      List[str]			=			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			:      Optional[int]			=			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			:      Dict			=			os.path.join(self.tmpdirname   ,  """psgs_w100.tsv.pkl"""		)
    UpperCAmelCase			:      Any			=			{sample["""id"""]: [sample["""text"""], sample["""title"""]] for sample in dataset}
    pickle.dump(_SCREAMING_SNAKE_CASE   ,  open(_SCREAMING_SNAKE_CASE   ,  """wb"""		)		)

    UpperCAmelCase			:      Dict			=			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			:      str			=			RagRetriever(
        _SCREAMING_SNAKE_CASE   ,  question_encoder_tokenizer=self.get_dpr_tokenizer()   ,  generator_tokenizer=self.get_bart_tokenizer()		)
    return retriever

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





    '''simple docstring'''




    UpperCAmelCase			:      int			=			1
    UpperCAmelCase			:      List[Any]			=			self.get_dummy_canonical_hf_index_retriever()
    UpperCAmelCase			:      List[str]			=			np.array(
        [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
    UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase			:      List[Any]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase			:      List[Any]			=			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			:      Tuple			=			RagRetriever.from_pretrained(_SCREAMING_SNAKE_CASE		)
      self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE		)
      UpperCAmelCase			:      List[Any]			=			np.array(
          [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
      UpperCAmelCase			:      Optional[Any]			=			retriever.retrieve(_SCREAMING_SNAKE_CASE   ,  n_docs=1		)
      self.assertTrue(out is not None		)

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[int]:





    '''simple docstring'''




    UpperCAmelCase			:      Tuple			=			1
    UpperCAmelCase			:      Dict			=			self.get_dummy_custom_hf_index_retriever(from_disk=_SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      Optional[Any]			=			np.array(
        [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
    UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase			:      Union[str, Any]			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Union[str, Any]:





    '''simple docstring'''




    UpperCAmelCase			:      Optional[int]			=			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			:      Union[str, Any]			=			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			:      List[str]			=			retriever.retrieve(_SCREAMING_SNAKE_CASE   ,  n_docs=1		)
     self.assertTrue(out is not None		)

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





    '''simple docstring'''




    UpperCAmelCase			:      Optional[Any]			=			1
    UpperCAmelCase			:      Optional[Any]			=			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			:      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 			SCREAMING_SNAKE_CASE  (     self		)						->   Optional[Any]:





    '''simple docstring'''




    UpperCAmelCase			:      Optional[int]			=			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			:      List[Any]			=			RagRetriever.from_pretrained(_SCREAMING_SNAKE_CASE		)
     self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE		)
     UpperCAmelCase			:      List[str]			=			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		)

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





    '''simple docstring'''




    UpperCAmelCase			:      Tuple			=			1
    UpperCAmelCase			:      List[Any]			=			self.get_dummy_legacy_index_retriever()
    UpperCAmelCase			:      Any			=			np.array(
        [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
    UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase			:      Dict			=			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 			SCREAMING_SNAKE_CASE  (     self		)						->   Tuple:





    '''simple docstring'''




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

   @require_torch
   @require_tokenizers
   @require_sentencepiece
   def 			SCREAMING_SNAKE_CASE  (     self		)						->   List[str]:





    '''simple docstring'''




    import torch

    UpperCAmelCase			:      Dict			=			1
    UpperCAmelCase			:      List[Any]			=			self.get_dummy_canonical_hf_index_retriever()
    UpperCAmelCase			:      List[str]			=			[[5, 7], [10, 11]]
    UpperCAmelCase			:      List[Any]			=			np.array(
        [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
    UpperCAmelCase			:      Union[str, Any]			=			retriever(_SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE   ,  prefix=retriever.config.generator.prefix   ,  n_docs=_SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase			:      Any			=			(
        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			:      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			:      Dict			=			(  # 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 			SCREAMING_SNAKE_CASE  (     self		)						->   Any:





    '''simple docstring'''




    UpperCAmelCase			:      int			=			self.get_dpr_ctx_encoder_tokenizer()
    UpperCAmelCase			:      Any			=			1
    UpperCAmelCase			:      List[str]			=			self.get_dummy_custom_hf_index_retriever(from_disk=_SCREAMING_SNAKE_CASE		)
    retriever.set_ctx_encoder_tokenizer(_SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Dict			=			[[5, 7], [10, 11]]
    UpperCAmelCase			:      Optional[Any]			=			np.array(
        [np.ones(self.retrieval_vector_size		), -np.ones(self.retrieval_vector_size		)]   ,  dtype=np.floataa		)
    UpperCAmelCase			:      Tuple			=			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.
