





"""simple docstring"""





import logging
import random

import ray

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


SCREAMING_SNAKE_CASE__ =       logging.getLogger(__name__)

class 		lowercase :






      """simple docstring"""




      def __init__(	self			)		->					str:
            lowerCAmelCase  						=	False





      def    _snake_case      (	self				,						lowercase				,						lowercase				,						lowercase				,						lowercase			)		->					List[str]:
            if not self.initialized:
                  lowerCAmelCase  						=	RagRetriever(
                      lowercase				,						question_encoder_tokenizer=lowercase				,						generator_tokenizer=lowercase				,						index=lowercase				,						init_retrieval=lowercase				,						)
                  lowerCAmelCase  						=	True





      def    _snake_case      (	self			)		->					Union[str, Any]:
            self.retriever.index.init_index()





      def    _snake_case      (	self			)		->					List[str]:
            # delete the old self.retriever object before assigning the new index
            del self.retriever
            lowerCAmelCase  						=	False





      def    _snake_case      (	self				,						lowercase				,						lowercase			)		->					Dict:
            lowerCAmelCase						,		lowerCAmelCase  						=	self.retriever._main_retrieve(lowercase				,						lowercase			)
            lowerCAmelCase  						=	self.retriever.index.get_doc_dicts(lowercase			)
            return doc_ids, retrieved_doc_embeds, doc_dicts







class 		lowercase ( _UpperCAmelCase ):






      """simple docstring"""




      def __init__(	self				,						lowercase				,						lowercase				,						lowercase				,						lowercase				,						lowercase=None			)		->					int:
            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				,						)

            lowerCAmelCase  						=	retrieval_workers
            lowerCAmelCase  						=	question_encoder_tokenizer
            lowerCAmelCase  						=	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    _snake_case      (	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    _snake_case      (	self				,						lowercase				,						lowercase			)		->					List[str]:
            if len(self.retrieval_workers			) > 0:
                  # Select a random retrieval actor.
                  lowerCAmelCase  						=	self.retrieval_workers[random.randint(0				,						len(self.retrieval_workers			) - 1			)]
                  lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=	ray.get(
                      random_worker.retrieve.remote(lowercase				,						lowercase			)			)
            else:
                  lowerCAmelCase						,		lowerCAmelCase  						=	self._main_retrieve(lowercase				,						lowercase			)
                  lowerCAmelCase  						=	self.index.get_doc_dicts(lowercase			)
            return retrieved_doc_embeds, doc_ids, doc_dicts





      @classmethod
      def    _snake_case      (	cls				,						lowercase				,						lowercase=None				,						**lowercase			)		->					Dict:
            return super(lowercase				,						cls			).get_tokenizers(lowercase				,						lowercase				,						**lowercase			)





      @classmethod
      def    _snake_case      (	cls				,						lowercase				,						lowercase				,						lowercase=None				,						**lowercase			)		->					str:
            lowerCAmelCase  						=	kwargs.pop("""config"""				,						lowercase			) or RagConfig.from_pretrained(lowercase				,						**lowercase			)
            lowerCAmelCase  						=	RagTokenizer.from_pretrained(lowercase				,						config=lowercase			)
            lowerCAmelCase  						=	rag_tokenizer.question_encoder
            lowerCAmelCase  						=	rag_tokenizer.generator

            if indexed_dataset is not None:
                  lowerCAmelCase  						=	"""custom"""
                  lowerCAmelCase  						=	CustomHFIndex(config.retrieval_vector_size				,						lowercase			)
            else:
                  lowerCAmelCase  						=	cls._build_index(lowercase			)

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





      def    _snake_case      (	self			)		->					List[str]:
            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
            lowerCAmelCase  						=	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
                ]			)













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

class 		lowercase ( unittest.TestCase ):





      def __init__(	self				,						lowercase				,						lowercase=7				,						lowercase=3				,						lowercase=18				,						lowercase=30				,						lowercase=400				,						lowercase=True				,						lowercase=None				,						lowercase=0.875				,						lowercase=True				,						lowercase=[0.5, 0.5, 0.5]				,						lowercase=[0.5, 0.5, 0.5]				,						)		->					List[Any]:
            lowerCAmelCase  						=	size if size is not None else {"""shortest_edge""": 20}
            lowerCAmelCase  						=	parent
            lowerCAmelCase  						=	batch_size
            lowerCAmelCase  						=	num_channels
            lowerCAmelCase  						=	image_size
            lowerCAmelCase  						=	min_resolution
            lowerCAmelCase  						=	max_resolution
            lowerCAmelCase  						=	do_resize
            lowerCAmelCase  						=	size
            lowerCAmelCase  						=	crop_pct
            lowerCAmelCase  						=	do_normalize
            lowerCAmelCase  						=	image_mean
            lowerCAmelCase  						=	image_std





      def    _snake_case      (	self			)		->					Optional[Any]:
            return {
                "image_mean": self.image_mean,
                "image_std": self.image_std,
                "do_normalize": self.do_normalize,
                "do_resize": self.do_resize,
                "size": self.size,
                "crop_pct": self.crop_pct,
            }







@require_torch
@require_vision
class 		lowercase ( _UpperCAmelCase					,  unittest.TestCase ):
      _SCREAMING_SNAKE_CASE         =				ConvNextImageProcessor if is_vision_available() else None





      def    _snake_case      (	self			)		->					List[str]:
            lowerCAmelCase  						=	ConvNextImageProcessingTester(self			)





      @property
      def    _snake_case      (	self			)		->					List[str]:
            return self.image_processor_tester.prepare_image_processor_dict()





      def    _snake_case      (	self			)		->					int:
            lowerCAmelCase  						=	self.image_processing_class(**self.image_processor_dict			)
            self.assertTrue(hasattr(lowercase				,						"""do_resize"""			)			)
            self.assertTrue(hasattr(lowercase				,						"""size"""			)			)
            self.assertTrue(hasattr(lowercase				,						"""crop_pct"""			)			)
            self.assertTrue(hasattr(lowercase				,						"""do_normalize"""			)			)
            self.assertTrue(hasattr(lowercase				,						"""image_mean"""			)			)
            self.assertTrue(hasattr(lowercase				,						"""image_std"""			)			)





      def    _snake_case      (	self			)		->					int:
            lowerCAmelCase  						=	self.image_processing_class.from_dict(self.image_processor_dict			)
            self.assertEqual(image_processor.size				,						{"""shortest_edge""": 20}			)

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





      def    _snake_case      (	self			)		->					Optional[int]:
            pass





      def    _snake_case      (	self			)		->					Optional[int]:
            # Initialize image_processing
            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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)

            # 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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)





      def    _snake_case      (	self			)		->					str:
            # Initialize image_processing
            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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)

            # 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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)





      def    _snake_case      (	self			)		->					Optional[int]:
            # Initialize image_processing
            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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)

            # 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.size["""shortest_edge"""],
                    self.image_processor_tester.size["""shortest_edge"""],
                )				,						)






