



"""simple docstring"""






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 ViltImageProcessor




class    __snake_case(					unittest.TestCase       ):




    def __init__(   self    ,							A_    ,							A_=7    ,							A_=3    ,							A_=18    ,							A_=30    ,							A_=400    ,							A_=True    ,							A_=None    ,							A_=2    ,							A_=True    ,							A_=[0.5, 0.5, 0.5]    ,							A_=[0.5, 0.5, 0.5]    ,							):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			size if size is not None else {'''shortest_edge''': 30}
      _SCREAMING_SNAKE_CASE        =			parent
      _SCREAMING_SNAKE_CASE        =			batch_size
      _SCREAMING_SNAKE_CASE        =			num_channels
      _SCREAMING_SNAKE_CASE        =			image_size
      _SCREAMING_SNAKE_CASE        =			min_resolution
      _SCREAMING_SNAKE_CASE        =			max_resolution
      _SCREAMING_SNAKE_CASE        =			do_resize
      _SCREAMING_SNAKE_CASE        =			size
      _SCREAMING_SNAKE_CASE        =			size_divisor
      _SCREAMING_SNAKE_CASE        =			do_normalize
      _SCREAMING_SNAKE_CASE        =			image_mean
      _SCREAMING_SNAKE_CASE        =			image_std




    def   A   (   self			):





      '''simple docstring'''



      return {
          "image_mean": self.image_mean,
          "image_std": self.image_std,
          "do_normalize": self.do_normalize,
          "do_resize": self.do_resize,
          "size": self.size,
          "size_divisor": self.size_divisor,
      }




    def   A   (   self    ,							A_    ,							A_=False			):





      '''simple docstring'''



      if not batched:
        _SCREAMING_SNAKE_CASE        =			self.size['''shortest_edge''']
        _SCREAMING_SNAKE_CASE        =			image_inputs[0]
        if isinstance(lowercase__    ,							Image.Image			):
          _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			image.size
        else:
          _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			image.shape[1], image.shape[2]
        _SCREAMING_SNAKE_CASE        =			size / min(lowercase__    ,							lowercase__			)
        if h < w:
          _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			size, scale * w
        else:
          _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			scale * h, size

        _SCREAMING_SNAKE_CASE        =			int((1_333 / 800) * size			)
        if max(lowercase__    ,							lowercase__			) > max_size:
          _SCREAMING_SNAKE_CASE        =			max_size / max(lowercase__    ,							lowercase__			)
          _SCREAMING_SNAKE_CASE        =			newh * scale
          _SCREAMING_SNAKE_CASE        =			neww * scale

        _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			int(newh + 0.5			), int(neww + 0.5			)
        _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			(
            newh // self.size_divisor * self.size_divisor,
            neww // self.size_divisor * self.size_divisor,
        )

      else:
        _SCREAMING_SNAKE_CASE        =			[]
        for image in image_inputs:
          _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.get_expected_values([image]			)
          expected_values.append((expected_height, expected_width)			)
        _SCREAMING_SNAKE_CASE        =			max(lowercase__    ,							key=lambda A_			: item[0]			)[0]
        _SCREAMING_SNAKE_CASE        =			max(lowercase__    ,							key=lambda A_			: item[1]			)[1]

      return expected_height, expected_width




@require_torch
@require_vision
class    __snake_case(					a__			,			unittest.TestCase       ):
    _A       = ViltImageProcessor if is_vision_available() else None




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			ViltImageProcessingTester(self			)




    @property
    def   A   (   self			):





      '''simple docstring'''



      return self.image_processor_tester.prepare_image_processor_dict()




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.image_processing_class(**self.image_processor_dict			)
      self.assertTrue(hasattr(lowercase__    ,							'''image_mean'''			)			)
      self.assertTrue(hasattr(lowercase__    ,							'''image_std'''			)			)
      self.assertTrue(hasattr(lowercase__    ,							'''do_normalize'''			)			)
      self.assertTrue(hasattr(lowercase__    ,							'''do_resize'''			)			)
      self.assertTrue(hasattr(lowercase__    ,							'''size'''			)			)
      self.assertTrue(hasattr(lowercase__    ,							'''size_divisor'''			)			)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.image_processing_class.from_dict(self.image_processor_dict			)
      self.assertEqual(image_processor.size    ,							{'''shortest_edge''': 30}			)

      _SCREAMING_SNAKE_CASE        =			self.image_processing_class.from_dict(self.image_processor_dict    ,							size=42			)
      self.assertEqual(image_processor.size    ,							{'''shortest_edge''': 42}			)




    def   A   (   self			):





      '''simple docstring'''



      pass




    def   A   (   self			):





      '''simple docstring'''



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

      # Test not batched input
      _SCREAMING_SNAKE_CASE        =			image_processing(image_inputs[0]    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,							)

      # Test batched
      _SCREAMING_SNAKE_CASE        =			image_processing(lowercase__    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__    ,							batched=lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )    ,							)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.image_processing_class(**self.image_processor_dict			)
      # create random numpy tensors
      _SCREAMING_SNAKE_CASE        =			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
      _SCREAMING_SNAKE_CASE        =			image_processing(image_inputs[0]    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,							)

      # Test batched
      _SCREAMING_SNAKE_CASE        =			image_processing(lowercase__    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__    ,							batched=lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )    ,							)




    def   A   (   self			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			self.image_processing_class(**self.image_processor_dict			)
      # create random PyTorch tensors
      _SCREAMING_SNAKE_CASE        =			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
      _SCREAMING_SNAKE_CASE        =			image_processing(image_inputs[0]    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(1, self.image_processor_tester.num_channels, expected_height, expected_width)    ,							)

      # Test batched
      _SCREAMING_SNAKE_CASE        =			image_processing(lowercase__    ,							return_tensors='''pt'''			).pixel_values

      _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			self.image_processor_tester.get_expected_values(lowercase__    ,							batched=lowercase__			)
      self.assertEqual(
          encoded_images.shape    ,							(
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )    ,							)






import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


UpperCamelCase__       				=	logging.getLogger(__name__)
class 							__lowercase   :



     def __init__(  self   :							List[str]       ):
       a_		    =  False



     def 				__magic_name__		(  self   :							Union[str, Any]      ,		lowercase__   :							List[str]      ,		lowercase__   :							Optional[int]      ,		lowercase__   :							Optional[Any]      ,		lowercase__   :							Union[str, Any]       ):
       if not self.initialized:
         a_		    =  RagRetriever(
             lowercase__      ,		question_encoder_tokenizer=lowercase__      ,		generator_tokenizer=lowercase__      ,		index=lowercase__      ,		init_retrieval=lowercase__      ,		)
         a_		    =  True



     def 				__magic_name__		(  self   :							List[Any]       ):
       self.retriever.index.init_index()



     def 				__magic_name__		(  self   :							List[Any]       ):
       # delete the old self.retriever object before assigning the new index
       del self.retriever
       a_		    =  False



     def 				__magic_name__		(  self   :							Optional[Any]      ,		lowercase__   :							List[str]      ,		lowercase__   :							Dict       ):
       a_    , a_		    =  self.retriever._main_retrieve(lowercase__      ,		lowercase__       )
       a_		    =  self.retriever.index.get_doc_dicts(lowercase__       )
       return doc_ids, retrieved_doc_embeds, doc_dicts



class 							__lowercase   (			a__    ):



     def __init__(  self   :							List[Any]      ,		lowercase__   :							Optional[int]      ,		lowercase__   :							List[str]      ,		lowercase__   :							Tuple      ,		lowercase__   :							Dict      ,		lowercase__   :							Dict=None       ):
       if index is not None and index.is_initialized() and len(lowercase__       ) > 0:
         raise ValueError(
             '''When using Ray for distributed fine-tuning, '''
             '''you\'ll need to provide the paths instead, '''
             '''as the dataset and the index are loaded '''
             '''separately. More info in examples/rag/use_own_knowledge_dataset.py '''       )

       super().__init__(
           lowercase__      ,		question_encoder_tokenizer=lowercase__      ,		generator_tokenizer=lowercase__      ,		index=lowercase__      ,		init_retrieval=lowercase__      ,		)

       a_		    =  retrieval_workers
       a_		    =  question_encoder_tokenizer
       a_		    =  generator_tokenizer
       if len(self.retrieval_workers       ) > 0:
         ray.get(
             [
                 worker.create_rag_retriever.remote(lowercase__      ,		lowercase__      ,		lowercase__      ,		lowercase__       )
                 for worker in self.retrieval_workers
             ]       )



     def 				__magic_name__		(  self   :							int       ):
       logger.info('''initializing retrieval'''       )

       if len(self.retrieval_workers       ) > 0:
         ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]       )
       else:
         # Non-distributed training. Load index into this same process.
         self.index.init_index()



     def 				__magic_name__		(  self   :							str      ,		lowercase__   :							List[str]      ,		lowercase__   :							Any       ):
       if len(self.retrieval_workers       ) > 0:
         # Select a random retrieval actor.
         a_		    =  self.retrieval_workers[random.randint(0      ,		len(self.retrieval_workers       ) - 1       )]
         a_    , a_    , a_		    =  ray.get(
             random_worker.retrieve.remote(lowercase__      ,		lowercase__       )       )
       else:
         a_    , a_		    =  self._main_retrieve(lowercase__      ,		lowercase__       )
         a_		    =  self.index.get_doc_dicts(lowercase__       )
       return retrieved_doc_embeds, doc_ids, doc_dicts



     @classmethod
     def 				__magic_name__		(  cls   :							Dict      ,		lowercase__   :							Optional[Any]      ,		lowercase__   :							List[str]=None      ,		**lowercase__   :							int       ):
       return super(lowercase__      ,		cls       ).get_tokenizers(lowercase__      ,		lowercase__      ,		**lowercase__       )



     @classmethod
     def 				__magic_name__		(  cls   :							Tuple      ,		lowercase__   :							Any      ,		lowercase__   :							str      ,		lowercase__   :							List[Any]=None      ,		**lowercase__   :							Union[str, Any]       ):
       a_		    =  kwargs.pop('''config'''      ,		lowercase__       ) or RagConfig.from_pretrained(lowercase__      ,		**lowercase__       )
       a_		    =  RagTokenizer.from_pretrained(lowercase__      ,		config=lowercase__       )
       a_		    =  rag_tokenizer.question_encoder
       a_		    =  rag_tokenizer.generator

       if indexed_dataset is not None:
         a_		    =  '''custom'''
         a_		    =  CustomHFIndex(config.retrieval_vector_size      ,		lowercase__       )
       else:
         a_		    =  cls._build_index(lowercase__       )

       return cls(
           lowercase__      ,		question_encoder_tokenizer=lowercase__      ,		generator_tokenizer=lowercase__      ,		retrieval_workers=lowercase__      ,		index=lowercase__      ,		)



     def 				__magic_name__		(  self   :							Optional[Any]       ):
       logger.info('''re-loading the new dataset with embeddings'''       )
       # access from the training loop

       ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]       )

       # build the index object again
       a_		    =  self._build_index(self.config       )

       ray.get(
           [
               worker.create_rag_retriever.remote(
                   self.config      ,		self.question_encoder_tokenizer      ,		self.generator_tokenizer      ,		lowercase__       )
               for worker in self.retrieval_workers
           ]       )
