




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
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 MobileNetVaImageProcessor

class 		lowercase_       (	unittest.TestCase				):




					def __init__(   self:		Dict,     _lowercase:		List[Any],     _lowercase:		int=7,     _lowercase:		List[str]=3,     _lowercase:		Optional[int]=18,     _lowercase:		Dict=30,     _lowercase:		Union[str, Any]=400,     _lowercase:		Union[str, Any]=True,     _lowercase:		Any=None,     _lowercase:		Optional[Any]=True,     _lowercase:		Any=None,     ):


									'''simple docstring'''


									__lowerCAmelCase					=  size if size is not None else {"""shortest_edge""": 20}
									__lowerCAmelCase					=  crop_size if crop_size is not None else {"""height""": 18, """width""": 18}
									__lowerCAmelCase					=  parent
									__lowerCAmelCase					=  batch_size
									__lowerCAmelCase					=  num_channels
									__lowerCAmelCase					=  image_size
									__lowerCAmelCase					=  min_resolution
									__lowerCAmelCase					=  max_resolution
									__lowerCAmelCase					=  do_resize
									__lowerCAmelCase					=  size
									__lowerCAmelCase					=  do_center_crop
									__lowerCAmelCase					=  crop_size







					def     _lowercase						(   self:		Dict):


									'''simple docstring'''


									return {
									    "do_resize": self.do_resize,
									    "size": self.size,
									    "do_center_crop": self.do_center_crop,
									    "crop_size": self.crop_size,
									}





@require_torch
@require_vision
class 		lowercase_       (	lowerCAmelCase__						,							unittest.TestCase				):
					__UpperCamelCase     	= MobileNetVaImageProcessor if is_vision_available() else None




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


									'''simple docstring'''


									__lowerCAmelCase					=  MobileNetVaImageProcessingTester(self)




					@property
					def     _lowercase						(   self:		Union[str, Any]):


									'''simple docstring'''


									return self.image_processor_tester.prepare_image_processor_dict()




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									__lowerCAmelCase					=  self.image_processing_class(**self.image_processor_dict)
									self.assertTrue(hasattr(_lowercase,     """do_resize"""))
									self.assertTrue(hasattr(_lowercase,     """size"""))
									self.assertTrue(hasattr(_lowercase,     """do_center_crop"""))
									self.assertTrue(hasattr(_lowercase,     """center_crop"""))




					def     _lowercase						(   self:		str):


									'''simple docstring'''


									__lowerCAmelCase					=  self.image_processing_class.from_dict(self.image_processor_dict)
									self.assertEqual(image_processor.size,     {"""shortest_edge""": 20})
									self.assertEqual(image_processor.crop_size,     {"""height""": 18, """width""": 18})

									__lowerCAmelCase					=  self.image_processing_class.from_dict(self.image_processor_dict,     size=42,     crop_size=84)
									self.assertEqual(image_processor.size,     {"""shortest_edge""": 42})
									self.assertEqual(image_processor.crop_size,     {"""height""": 84, """width""": 84})




					def     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


									pass




					def     _lowercase						(   self:		List[Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  self.image_processing_class(**self.image_processor_dict)
									# create random PIL images
									__lowerCAmelCase					=  prepare_image_inputs(self.image_processor_tester,     equal_resolution=_lowercase)
									for image in image_inputs:
													self.assertIsInstance(_lowercase,     Image.Image)

									# Test not batched input
									__lowerCAmelCase					=  image_processing(image_inputs[0],     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )

									# Test batched
									__lowerCAmelCase					=  image_processing(_lowercase,     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )




					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


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

									# Test not batched input
									__lowerCAmelCase					=  image_processing(image_inputs[0],     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )

									# Test batched
									__lowerCAmelCase					=  image_processing(_lowercase,     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )







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


									'''simple docstring'''


									__lowerCAmelCase					=  self.image_processing_class(**self.image_processor_dict)
									# create random PyTorch tensors
									__lowerCAmelCase					=  prepare_image_inputs(self.image_processor_tester,     equal_resolution=_lowercase,     torchify=_lowercase)
									for image in image_inputs:
													self.assertIsInstance(_lowercase,     torch.Tensor)

									# Test not batched input
									__lowerCAmelCase					=  image_processing(image_inputs[0],     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        1,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )

									# Test batched
									__lowerCAmelCase					=  image_processing(_lowercase,     return_tensors="""pt""").pixel_values
									self.assertEqual(
									    encoded_images.shape,     (
									        self.image_processor_tester.batch_size,
									        self.image_processor_tester.num_channels,
									        self.image_processor_tester.crop_size["""height"""],
									        self.image_processor_tester.crop_size["""width"""],
									    ),     )











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

import faiss
import numpy as np
from datasets import Dataset

from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
from transformers.integrations import is_ray_available
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
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_ray


sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

if is_torch_available():
	from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
else:
	__A						:		str          =						None

if is_ray_available():
	import ray  # noqa: E402 # isort:skip
	from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
else:
	__A						:		int          =						None
	__A						:		Tuple          =						None
	__A						:		Union[str, Any]          =						None





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







				'''simple docstring'''


				if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
								__lowerCAmelCase					=  unittest.skip("""test requires Datasets, Faiss, psutil"""   )(UpperCamelCase__   )
				return test_case

@require_distributed_retrieval
class 		lowercase_       (	lowerCAmelCase__				):




					def     _lowercase						(   self:		List[Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  tempfile.mkdtemp()
									__lowerCAmelCase					=  8

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

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




					def     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


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




					def     _lowercase						(   self:		int):


									'''simple docstring'''


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




					def     _lowercase						(   self:		Optional[int]):


									'''simple docstring'''


									shutil.rmtree(self.tmpdirname)




					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


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




					def     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool,     _lowercase:		Any=12345):


									'''simple docstring'''


									__lowerCAmelCase					=  self.get_dummy_dataset()
									__lowerCAmelCase					=  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:
													__lowerCAmelCase					=  dataset
													__lowerCAmelCase					=  RagPyTorchDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
													if init_retrieval:
																	retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		Union[str, Any],     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__lowerCAmelCase					=  RagConfig(
									    retrieval_vector_size=self.retrieval_vector_size,     question_encoder=DPRConfig().to_dict(),     generator=BartConfig().to_dict(),     )
									__lowerCAmelCase					=  ray.remote(_lowercase)
									__lowerCAmelCase					=  [remote_cls.remote() for _ in range(1)]
									with patch("""transformers.models.rag.retrieval_rag.load_dataset""") as mock_load_dataset:
													__lowerCAmelCase					=  self.get_dummy_dataset()
													__lowerCAmelCase					=  RagRayDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     )
													if init_retrieval:
																	retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		Optional[Any],     _lowercase:		bool,     _lowercase:		bool,     _lowercase:		int=12345):


									'''simple docstring'''


									__lowerCAmelCase					=  self.get_dummy_dataset()
									__lowerCAmelCase					=  RagConfig(
									    retrieval_vector_size=self.retrieval_vector_size,     question_encoder=DPRConfig().to_dict(),     generator=BartConfig().to_dict(),     index_name="""custom""",     )
									if from_disk:
													__lowerCAmelCase					=  os.path.join(self.tmpdirname,     """dataset""")
													__lowerCAmelCase					=  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
													__lowerCAmelCase					=  RagPyTorchDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     )
									else:
													__lowerCAmelCase					=  RagPyTorchDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval(_lowercase)
									return retriever




					def     _lowercase						(   self:		List[Any],     _lowercase:		bool,     _lowercase:		bool):


									'''simple docstring'''


									ray.init(local_mode=_lowercase)
									__lowerCAmelCase					=  self.get_dummy_dataset()
									__lowerCAmelCase					=  RagConfig(
									    retrieval_vector_size=self.retrieval_vector_size,     question_encoder=DPRConfig().to_dict(),     generator=BartConfig().to_dict(),     index_name="""custom""",     )
									__lowerCAmelCase					=  ray.remote(_lowercase)
									__lowerCAmelCase					=  [remote_cls.remote() for _ in range(1)]
									if from_disk:
													__lowerCAmelCase					=  os.path.join(self.tmpdirname,     """dataset""")
													__lowerCAmelCase					=  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
													__lowerCAmelCase					=  RagRayDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex.load_from_disk(
													        vector_size=config.retrieval_vector_size,     dataset_path=config.passages_path,     index_path=config.index_path,     ),     )
									else:
													__lowerCAmelCase					=  RagRayDistributedRetriever(
													    _lowercase,     question_encoder_tokenizer=self.get_dpr_tokenizer(),     generator_tokenizer=self.get_bart_tokenizer(),     retrieval_workers=_lowercase,     index=CustomHFIndex(config.retrieval_vector_size,     _lowercase),     )
									if init_retrieval:
													retriever.init_retrieval()
									return retriever




					def     _lowercase						(   self:		int,     _lowercase:		RagRetriever,     _lowercase:		np.array,     _lowercase:		int):


									'''simple docstring'''


									__lowerCAmelCase	,				__lowerCAmelCase	,				__lowerCAmelCase					=  retriever.retrieve(_lowercase,     n_docs=_lowercase)
									self.assertEqual(retrieved_doc_embeds.shape,     (2, n_docs, self.retrieval_vector_size))
									self.assertEqual(len(_lowercase),     2)
									self.assertEqual(sorted(doc_dicts[0]),     ["""embeddings""", """id""", """text""", """title"""])
									self.assertEqual(len(doc_dicts[0]["""id"""]),     _lowercase)
									self.assertEqual(doc_dicts[0]["""id"""][0],     """1""")  # max inner product is reached with second doc
									self.assertEqual(doc_dicts[1]["""id"""][0],     """0""")  # max inner product is reached with first doc
									self.assertListEqual(doc_ids.tolist(),     [[1], [0]])




					def     _lowercase						(   self:		Optional[Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)

									self.distributed_retriever_check(
									    self.get_dummy_pytorch_distributed_retriever(init_retrieval=_lowercase),     _lowercase,     _lowercase)




					def     _lowercase						(   self:		int):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)

									self.distributed_retriever_check(
									    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					def     _lowercase						(   self:		Dict):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)

									self.distributed_retriever_check(
									    self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )




					@require_ray
					def     _lowercase						(   self:		List[str]):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)

									self.distributed_retriever_check(
									    self.get_dummy_ray_distributed_retriever(init_retrieval=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()




					@require_ray
					def     _lowercase						(   self:		Union[str, Any]):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)
									with self.assertRaises(_lowercase):
													self.distributed_retriever_check(
													    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase,     )
									ray.shutdown()







					@require_ray
					def     _lowercase						(   self:		Tuple):


									'''simple docstring'''


									__lowerCAmelCase					=  1
									__lowerCAmelCase					=  np.array(
									    [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)],     dtype=np.floataa)

									self.distributed_retriever_check(
									    self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=_lowercase,     from_disk=_lowercase),     _lowercase,     _lowercase)
									ray.shutdown()





