



"""simple docstring"""
import json
import pathlib
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision, slow
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import YolosImageProcessor



class    SCREAMING_SNAKE_CASE__						(      unittest.TestCase				):

   def __init__(     self   ,  _SCREAMING_SNAKE_CASE   ,  _SCREAMING_SNAKE_CASE=7   ,  _SCREAMING_SNAKE_CASE=3   ,  _SCREAMING_SNAKE_CASE=30   ,  _SCREAMING_SNAKE_CASE=400   ,  _SCREAMING_SNAKE_CASE=True   ,  _SCREAMING_SNAKE_CASE=None   ,  _SCREAMING_SNAKE_CASE=True   ,  _SCREAMING_SNAKE_CASE=[0.5, 0.5, 0.5]   ,  _SCREAMING_SNAKE_CASE=[0.5, 0.5, 0.5]   ,  _SCREAMING_SNAKE_CASE=True   ,  _SCREAMING_SNAKE_CASE=1 / 255   ,  _SCREAMING_SNAKE_CASE=True   ,  )						->   Optional[int]:





    '''simple docstring'''




    UpperCAmelCase			:      List[Any]			=			size if size is not None else {"""shortest_edge""": 18, """longest_edge""": 1333}
    UpperCAmelCase			:      Optional[int]			=			parent
    UpperCAmelCase			:      Tuple			=			batch_size
    UpperCAmelCase			:      Optional[Any]			=			num_channels
    UpperCAmelCase			:      List[Any]			=			min_resolution
    UpperCAmelCase			:      Optional[Any]			=			max_resolution
    UpperCAmelCase			:      Dict			=			do_resize
    UpperCAmelCase			:      Optional[int]			=			size
    UpperCAmelCase			:      Optional[int]			=			do_normalize
    UpperCAmelCase			:      Tuple			=			image_mean
    UpperCAmelCase			:      int			=			image_std
    UpperCAmelCase			:      List[str]			=			do_rescale
    UpperCAmelCase			:      Dict			=			rescale_factor
    UpperCAmelCase			:      int			=			do_pad

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   int:





    '''simple docstring'''




    return {
        "do_resize": self.do_resize,
        "size": self.size,
        "do_normalize": self.do_normalize,
        "image_mean": self.image_mean,
        "image_std": self.image_std,
        "do_rescale": self.do_rescale,
        "rescale_factor": self.rescale_factor,
        "do_pad": self.do_pad,
    }




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





    '''simple docstring'''




    if not batched:
     UpperCAmelCase			:      Union[str, Any]			=			image_inputs[0]
     if isinstance(_SCREAMING_SNAKE_CASE   ,  Image.Image		):
      UpperCAmelCase			,	UpperCAmelCase			:      str			=			image.size
     else:
      UpperCAmelCase			,	UpperCAmelCase			:      Dict			=			image.shape[1], image.shape[2]
     if w < h:
      UpperCAmelCase			:      str			=			int(self.size["""shortest_edge"""] * h / w		)
      UpperCAmelCase			:      Optional[Any]			=			self.size["""shortest_edge"""]
     elif w > h:
      UpperCAmelCase			:      Any			=			self.size["""shortest_edge"""]
      UpperCAmelCase			:      Union[str, Any]			=			int(self.size["""shortest_edge"""] * w / h		)
     else:
      UpperCAmelCase			:      Tuple			=			self.size["""shortest_edge"""]
      UpperCAmelCase			:      Optional[int]			=			self.size["""shortest_edge"""]

    else:
     UpperCAmelCase			:      str			=			[]
     for image in image_inputs:
      UpperCAmelCase			,	UpperCAmelCase			:      List[Any]			=			self.get_expected_values([image]		)
      expected_values.append((expected_height, expected_width)		)
     UpperCAmelCase			:      List[str]			=			max(_SCREAMING_SNAKE_CASE   ,  key=lambda _SCREAMING_SNAKE_CASE		: item[0]		)[0]
     UpperCAmelCase			:      Tuple			=			max(_SCREAMING_SNAKE_CASE   ,  key=lambda _SCREAMING_SNAKE_CASE		: item[1]		)[1]

    return expected_height, expected_width



@require_torch
@require_vision
class    SCREAMING_SNAKE_CASE__						(      UpperCAmelCase__   ,    unittest.TestCase				):
   __lowerCAmelCase						:     Any		=		YolosImageProcessor if is_vision_available() else None

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





    '''simple docstring'''




    UpperCAmelCase			:      List[str]			=			YolosImageProcessingTester(self		)

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





    '''simple docstring'''




    return self.image_processor_tester.prepare_image_processor_dict()

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   str:





    '''simple docstring'''




    UpperCAmelCase			:      Dict			=			self.image_processing_class(**self.image_processor_dict		)
    self.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,  """image_mean"""		)		)
    self.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,  """image_std"""		)		)
    self.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,  """do_normalize"""		)		)
    self.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,  """do_resize"""		)		)
    self.assertTrue(hasattr(_SCREAMING_SNAKE_CASE   ,  """size"""		)		)

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      str			=			self.image_processing_class.from_dict(self.image_processor_dict		)
    self.assertEqual(image_processor.size   ,  {"""shortest_edge""": 18, """longest_edge""": 1333}		)
    self.assertEqual(image_processor.do_pad   ,  _SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Dict			=			self.image_processing_class.from_dict(
        self.image_processor_dict   ,  size=42   ,  max_size=84   ,  pad_and_return_pixel_mask=_SCREAMING_SNAKE_CASE		)
    self.assertEqual(image_processor.size   ,  {"""shortest_edge""": 42, """longest_edge""": 84}		)
    self.assertEqual(image_processor.do_pad   ,  _SCREAMING_SNAKE_CASE		)

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





    '''simple docstring'''




    pass

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      List[str]			=			self.image_processing_class(**self.image_processor_dict		)
    # create random PIL images
    UpperCAmelCase			:      Optional[int]			=			prepare_image_inputs(self.image_processor_tester   ,  equal_resolution=_SCREAMING_SNAKE_CASE		)
    for image in image_inputs:
     self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  Image.Image		)

    # Test not batched input
    UpperCAmelCase			:      List[str]			=			image_processing(image_inputs[0]   ,  return_tensors="""pt"""		).pixel_values

    UpperCAmelCase			,	UpperCAmelCase			:      List[Any]			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE		)

    self.assertEqual(
        encoded_images.shape   ,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)   ,  )

    # Test batched
    UpperCAmelCase			,	UpperCAmelCase			:      Optional[int]			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE   ,  batched=_SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      List[str]			=			image_processing(_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		).pixel_values
    self.assertEqual(
        encoded_images.shape   ,  (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            expected_height,
            expected_width,
        )   ,  )

   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      str			=			self.image_processing_class(**self.image_processor_dict		)
    # create random numpy tensors
    UpperCAmelCase			:      Any			=			prepare_image_inputs(self.image_processor_tester   ,  equal_resolution=_SCREAMING_SNAKE_CASE   ,  numpify=_SCREAMING_SNAKE_CASE		)
    for image in image_inputs:
     self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  np.ndarray		)

    # Test not batched input
    UpperCAmelCase			:      Optional[Any]			=			image_processing(image_inputs[0]   ,  return_tensors="""pt"""		).pixel_values

    UpperCAmelCase			,	UpperCAmelCase			:      Union[str, Any]			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE		)

    self.assertEqual(
        encoded_images.shape   ,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)   ,  )

    # Test batched
    UpperCAmelCase			:      Tuple			=			image_processing(_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		).pixel_values

    UpperCAmelCase			,	UpperCAmelCase			:      List[Any]			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE   ,  batched=_SCREAMING_SNAKE_CASE		)

    self.assertEqual(
        encoded_images.shape   ,  (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            expected_height,
            expected_width,
        )   ,  )

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





    '''simple docstring'''




    UpperCAmelCase			:      Optional[int]			=			self.image_processing_class(**self.image_processor_dict		)
    # create random PyTorch tensors
    UpperCAmelCase			:      List[str]			=			prepare_image_inputs(self.image_processor_tester   ,  equal_resolution=_SCREAMING_SNAKE_CASE   ,  torchify=_SCREAMING_SNAKE_CASE		)
    for image in image_inputs:
     self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  torch.Tensor		)

    # Test not batched input
    UpperCAmelCase			:      Any			=			image_processing(image_inputs[0]   ,  return_tensors="""pt"""		).pixel_values

    UpperCAmelCase			,	UpperCAmelCase			:      Any			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE		)

    self.assertEqual(
        encoded_images.shape   ,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)   ,  )

    # Test batched
    UpperCAmelCase			:      Dict			=			image_processing(_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		).pixel_values

    UpperCAmelCase			,	UpperCAmelCase			:      Union[str, Any]			=			self.image_processor_tester.get_expected_values(_SCREAMING_SNAKE_CASE   ,  batched=_SCREAMING_SNAKE_CASE		)

    self.assertEqual(
        encoded_images.shape   ,  (
            self.image_processor_tester.batch_size,
            self.image_processor_tester.num_channels,
            expected_height,
            expected_width,
        )   ,  )

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





    '''simple docstring'''




    UpperCAmelCase			:      int			=			self.image_processing_class(**self.image_processor_dict		)
    UpperCAmelCase			:      str			=			self.image_processing_class(do_resize=_SCREAMING_SNAKE_CASE   ,  do_normalize=_SCREAMING_SNAKE_CASE   ,  do_rescale=_SCREAMING_SNAKE_CASE		)
    # create random PyTorch tensors
    UpperCAmelCase			:      Tuple			=			prepare_image_inputs(self.image_processor_tester   ,  equal_resolution=_SCREAMING_SNAKE_CASE   ,  torchify=_SCREAMING_SNAKE_CASE		)
    for image in image_inputs:
     self.assertIsInstance(_SCREAMING_SNAKE_CASE   ,  torch.Tensor		)

    # Test whether the method "pad" and calling the image processor return the same tensors
    UpperCAmelCase			:      int			=			image_processing_a.pad(_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		)
    UpperCAmelCase			:      List[Any]			=			image_processing_a(_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		)

    self.assertTrue(
        torch.allclose(encoded_images_with_method["""pixel_values"""]   ,  encoded_images["""pixel_values"""]   ,  atol=1E-4		)		)

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





    '''simple docstring'''




    UpperCAmelCase			:      Optional[Any]			=			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
    with open("""./tests/fixtures/tests_samples/COCO/coco_annotations.txt"""   ,  """r"""		) as f:
     UpperCAmelCase			:      Optional[Any]			=			json.loads(f.read()		)

    UpperCAmelCase			:      Union[str, Any]			=			{"""image_id""": 39769, """annotations""": target}

    # encode them
    UpperCAmelCase			:      Optional[Any]			=			YolosImageProcessor.from_pretrained("""hustvl/yolos-small"""		)
    UpperCAmelCase			:      Dict			=			image_processing(images=_SCREAMING_SNAKE_CASE   ,  annotations=_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		)

    # verify pixel values
    UpperCAmelCase			:      Any			=			torch.Size([1, 3, 800, 1066]		)
    self.assertEqual(encoding["""pixel_values"""].shape   ,  _SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Optional[Any]			=			torch.tensor([0.2796, 0.3138, 0.3481]		)
    self.assertTrue(torch.allclose(encoding["""pixel_values"""][0, 0, 0, :3]   ,  _SCREAMING_SNAKE_CASE   ,  atol=1E-4		)		)

    # verify area
    UpperCAmelCase			:      str			=			torch.tensor([5887.9600, 1_1250.2061, 48_9353.8438, 83_7122.7500, 14_7967.5156, 16_5732.3438]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""area"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify boxes
    UpperCAmelCase			:      Dict			=			torch.Size([6, 4]		)
    self.assertEqual(encoding["""labels"""][0]["""boxes"""].shape   ,  _SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      int			=			torch.tensor([0.5503, 0.2765, 0.0604, 0.2215]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""boxes"""][0]   ,  _SCREAMING_SNAKE_CASE   ,  atol=1E-3		)		)
    # verify image_id
    UpperCAmelCase			:      Optional[int]			=			torch.tensor([39769]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""image_id"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify is_crowd
    UpperCAmelCase			:      Optional[int]			=			torch.tensor([0, 0, 0, 0, 0, 0]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""iscrowd"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify class_labels
    UpperCAmelCase			:      int			=			torch.tensor([75, 75, 63, 65, 17, 17]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""class_labels"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify orig_size
    UpperCAmelCase			:      str			=			torch.tensor([480, 640]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""orig_size"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify size
    UpperCAmelCase			:      int			=			torch.tensor([800, 1066]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""size"""]   ,  _SCREAMING_SNAKE_CASE		)		)




   @slow
   def 			SCREAMING_SNAKE_CASE  (     self		)						->   Dict:





    '''simple docstring'''




    UpperCAmelCase			:      List[str]			=			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""		)
    with open("""./tests/fixtures/tests_samples/COCO/coco_panoptic_annotations.txt"""   ,  """r"""		) as f:
     UpperCAmelCase			:      int			=			json.loads(f.read()		)

    UpperCAmelCase			:      Optional[Any]			=			{"""file_name""": """000000039769.png""", """image_id""": 39769, """segments_info""": target}

    UpperCAmelCase			:      Any			=			pathlib.Path("""./tests/fixtures/tests_samples/COCO/coco_panoptic"""		)

    # encode them
    UpperCAmelCase			:      Optional[int]			=			YolosImageProcessor(format="""coco_panoptic"""		)
    UpperCAmelCase			:      Any			=			image_processing(images=_SCREAMING_SNAKE_CASE   ,  annotations=_SCREAMING_SNAKE_CASE   ,  masks_path=_SCREAMING_SNAKE_CASE   ,  return_tensors="""pt"""		)

    # verify pixel values
    UpperCAmelCase			:      int			=			torch.Size([1, 3, 800, 1066]		)
    self.assertEqual(encoding["""pixel_values"""].shape   ,  _SCREAMING_SNAKE_CASE		)

    UpperCAmelCase			:      Optional[int]			=			torch.tensor([0.2796, 0.3138, 0.3481]		)
    self.assertTrue(torch.allclose(encoding["""pixel_values"""][0, 0, 0, :3]   ,  _SCREAMING_SNAKE_CASE   ,  atol=1E-4		)		)

    # verify area
    UpperCAmelCase			:      Optional[Any]			=			torch.tensor([14_7979.6875, 16_5527.0469, 48_4638.5938, 1_1292.9375, 5879.6562, 7634.1147]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""area"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify boxes
    UpperCAmelCase			:      Union[str, Any]			=			torch.Size([6, 4]		)
    self.assertEqual(encoding["""labels"""][0]["""boxes"""].shape   ,  _SCREAMING_SNAKE_CASE		)
    UpperCAmelCase			:      str			=			torch.tensor([0.2625, 0.5437, 0.4688, 0.8625]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""boxes"""][0]   ,  _SCREAMING_SNAKE_CASE   ,  atol=1E-3		)		)
    # verify image_id
    UpperCAmelCase			:      int			=			torch.tensor([39769]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""image_id"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify is_crowd
    UpperCAmelCase			:      List[str]			=			torch.tensor([0, 0, 0, 0, 0, 0]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""iscrowd"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify class_labels
    UpperCAmelCase			:      Optional[Any]			=			torch.tensor([17, 17, 63, 75, 75, 93]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""class_labels"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify masks
    UpperCAmelCase			:      Optional[int]			=			822873
    self.assertEqual(encoding["""labels"""][0]["""masks"""].sum().item()   ,  _SCREAMING_SNAKE_CASE		)
    # verify orig_size
    UpperCAmelCase			:      str			=			torch.tensor([480, 640]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""orig_size"""]   ,  _SCREAMING_SNAKE_CASE		)		)
    # verify size
    UpperCAmelCase			:      Union[str, Any]			=			torch.tensor([800, 1066]		)
    self.assertTrue(torch.allclose(encoding["""labels"""][0]["""size"""]   ,  _SCREAMING_SNAKE_CASE		)		)





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