
"""simple docstring"""

import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
					from transformers.models.mobilevitva.modeling_mobilevitva import (
					    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
					    make_divisible,
					)


if is_vision_available():
					from PIL import Image

					from transformers import MobileViTImageProcessor






class 	_lowercase       (    __UpperCAmelCase		):




	def       _UpperCamelCase   (			self )		->							List[Any]:
			lowerCamelCase					:      List[Any]          =   self.config_class(**self.inputs_dict )
			self.parent.assertTrue(hasattr(UpperCAmelCase_		,			'width_multiplier' ) )






class 	_lowercase       :




	def __init__(			self		,			UpperCAmelCase_		,			UpperCAmelCase_=13		,			UpperCAmelCase_=64		,			UpperCAmelCase_=2		,			UpperCAmelCase_=3		,			UpperCAmelCase_="swish"		,			UpperCAmelCase_=3		,			UpperCAmelCase_=32		,			UpperCAmelCase_=0.1		,			UpperCAmelCase_=0.02		,			UpperCAmelCase_=True		,			UpperCAmelCase_=True		,			UpperCAmelCase_=10		,			UpperCAmelCase_=None		,			UpperCAmelCase_=0.25		,			UpperCAmelCase_=0.0		,			UpperCAmelCase_=0.0		,			)		->							Optional[int]:
			lowerCamelCase					:      Tuple          =   parent
			lowerCamelCase					:      List[str]          =   batch_size
			lowerCamelCase					:      int          =   image_size
			lowerCamelCase					:      Tuple          =   patch_size
			lowerCamelCase					:      Optional[int]          =   num_channels
			lowerCamelCase					:      Union[str, Any]          =   make_divisible(512 * width_multiplier		,			divisor=8 )
			lowerCamelCase					:      Optional[Any]          =   hidden_act
			lowerCamelCase					:      Optional[Any]          =   conv_kernel_size
			lowerCamelCase					:      Any          =   output_stride
			lowerCamelCase					:      Tuple          =   classifier_dropout_prob
			lowerCamelCase					:      Optional[int]          =   use_labels
			lowerCamelCase					:      Any          =   is_training
			lowerCamelCase					:      Union[str, Any]          =   num_labels
			lowerCamelCase					:      int          =   initializer_range
			lowerCamelCase					:      List[str]          =   scope
			lowerCamelCase					:      List[str]          =   width_multiplier
			lowerCamelCase					:      Tuple          =   ffn_dropout
			lowerCamelCase					:      Union[str, Any]          =   attn_dropout




	def       _UpperCamelCase   (			self )		->							Dict:
			lowerCamelCase					:      Union[str, Any]          =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

			lowerCamelCase					:      Union[str, Any]          =   None
			lowerCamelCase					:      List[Any]          =   None
			if self.use_labels:
					lowerCamelCase					:      Dict          =   ids_tensor([self.batch_size]		,			self.num_labels )
					lowerCamelCase					:      List[Any]          =   ids_tensor([self.batch_size, self.image_size, self.image_size]		,			self.num_labels )

			lowerCamelCase					:      Any          =   self.get_config()

			return config, pixel_values, labels, pixel_labels




	def       _UpperCamelCase   (			self )		->							str:
			return MobileViTVaConfig(
			    image_size=self.image_size		,			patch_size=self.patch_size		,			num_channels=self.num_channels		,			hidden_act=self.hidden_act		,			conv_kernel_size=self.conv_kernel_size		,			output_stride=self.output_stride		,			classifier_dropout_prob=self.classifier_dropout_prob		,			initializer_range=self.initializer_range		,			width_multiplier=self.width_multiplier		,			ffn_dropout=self.ffn_dropout_prob		,			attn_dropout=self.attn_dropout_prob		,			)




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							int:
			lowerCamelCase					:      str          =   MobileViTVaModel(config=UpperCAmelCase_ )
			model.to(UpperCAmelCase_ )
			model.eval()
			lowerCamelCase					:      Any          =   model(UpperCAmelCase_ )
			self.parent.assertEqual(
			    result.last_hidden_state.shape		,			(
			        self.batch_size,
			        self.last_hidden_size,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )		,			)




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Optional[int]:
			lowerCamelCase					:      Optional[Any]          =   self.num_labels
			lowerCamelCase					:      Optional[Any]          =   MobileViTVaForImageClassification(UpperCAmelCase_ )
			model.to(UpperCAmelCase_ )
			model.eval()
			lowerCamelCase					:      List[str]          =   model(UpperCAmelCase_		,			labels=UpperCAmelCase_ )
			self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels) )




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							str:
			lowerCamelCase					:      List[Any]          =   self.num_labels
			lowerCamelCase					:      str          =   MobileViTVaForSemanticSegmentation(UpperCAmelCase_ )
			model.to(UpperCAmelCase_ )
			model.eval()
			lowerCamelCase					:      Optional[Any]          =   model(UpperCAmelCase_ )
			self.parent.assertEqual(
			    result.logits.shape		,			(
			        self.batch_size,
			        self.num_labels,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )		,			)
			lowerCamelCase					:      Dict          =   model(UpperCAmelCase_		,			labels=UpperCAmelCase_ )
			self.parent.assertEqual(
			    result.logits.shape		,			(
			        self.batch_size,
			        self.num_labels,
			        self.image_size // self.output_stride,
			        self.image_size // self.output_stride,
			    )		,			)




	def       _UpperCamelCase   (			self )		->							Dict:
			lowerCamelCase					:      Optional[Any]          =   self.prepare_config_and_inputs()
			lowerCamelCase     ,   lowerCamelCase     ,   lowerCamelCase     ,   lowerCamelCase					:      int          =   config_and_inputs
			lowerCamelCase					:      Dict          =   {'pixel_values': pixel_values}
			return config, inputs_dict






@require_torch
class 	_lowercase       (    __UpperCAmelCase    ,	__UpperCAmelCase    ,	unittest.TestCase		):
	lowercase_				=     (
	    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
	    if is_torch_available()
	    else ()
	)

	lowercase_				=     (
	    {
	        'feature-extraction': MobileViTVaModel,
	        'image-classification': MobileViTVaForImageClassification,
	        'image-segmentation': MobileViTVaForSemanticSegmentation,
	    }
	    if is_torch_available()
	    else {}
	)

	lowercase_				=     False
	lowercase_				=     False
	lowercase_				=     False
	lowercase_				=     False




	def       _UpperCamelCase   (			self )		->							List[str]:
			lowerCamelCase					:      List[str]          =   MobileViTVaModelTester(self )
			lowerCamelCase					:      List[Any]          =   MobileViTVaConfigTester(self		,			config_class=UpperCAmelCase_		,			has_text_modality=UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							Tuple:
			self.config_tester.run_common_tests()




	@unittest.skip(reason='MobileViTV2 does not use inputs_embeds' )
	def       _UpperCamelCase   (			self )		->							Any:
			pass




	@unittest.skip(reason='MobileViTV2 does not support input and output embeddings' )
	def       _UpperCamelCase   (			self )		->							Dict:
			pass




	@unittest.skip(reason='MobileViTV2 does not output attentions' )
	def       _UpperCamelCase   (			self )		->							int:
			pass




	@require_torch_multi_gpu
	@unittest.skip(reason='Got `CUDA error: misaligned address` for tests after this one being run.' )
	def       _UpperCamelCase   (			self )		->							List[Any]:
			pass




	@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.' )
	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			pass




	def       _UpperCamelCase   (			self )		->							int:
			lowerCamelCase     ,   lowerCamelCase					:      str          =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					lowerCamelCase					:      int          =   model_class(UpperCAmelCase_ )
					lowerCamelCase					:      Dict          =   inspect.signature(model.forward )
					# signature.parameters is an OrderedDict => so arg_names order is deterministic
					lowerCamelCase					:      Tuple          =   [*signature.parameters.keys()]

					lowerCamelCase					:      Dict          =   ['pixel_values']
					self.assertListEqual(arg_names[:1]		,			UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							Any:
			lowerCamelCase					:      Dict          =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							Optional[int]:
			def check_hidden_states_output(UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ ):
					lowerCamelCase					:      Optional[int]          =   model_class(UpperCAmelCase_ )
					model.to(UpperCAmelCase_ )
					model.eval()

					with torch.no_grad():
							lowerCamelCase					:      Tuple          =   model(**self._prepare_for_class(UpperCAmelCase_		,			UpperCAmelCase_ ) )

					lowerCamelCase					:      Tuple          =   outputs.hidden_states

					lowerCamelCase					:      Optional[int]          =   5
					self.assertEqual(len(UpperCAmelCase_ )		,			UpperCAmelCase_ )

					# MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
					# with the width and height being successively divided by 2.
					lowerCamelCase					:      Optional[int]          =   2
					for i in range(len(UpperCAmelCase_ ) ):
							self.assertListEqual(
							    list(hidden_states[i].shape[-2:] )		,			[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]		,			)
							divisor *= 2

					self.assertEqual(self.model_tester.output_stride		,			divisor // 2 )

			lowerCamelCase     ,   lowerCamelCase					:      Optional[Any]          =   self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					lowerCamelCase					:      Optional[int]          =   True
					check_hidden_states_output(UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )

					# check that output_hidden_states also work using config
					del inputs_dict["output_hidden_states"]
					lowerCamelCase					:      str          =   True

					check_hidden_states_output(UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							List[str]:
			lowerCamelCase					:      int          =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							Optional[Any]:
			lowerCamelCase					:      List[Any]          =   self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_ )




	@slow
	def       _UpperCamelCase   (			self )		->							Dict:
			for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase					:      int          =   MobileViTVaModel.from_pretrained(UpperCAmelCase_ )
					self.assertIsNotNone(UpperCAmelCase_ )


def 		UpperCAmelCase (		):





		'''simple docstring'''


		lowerCamelCase					:      int          =   Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
		return image






@require_torch
@require_vision
class 	_lowercase       (    unittest.TestCase		):




	@cached_property
	def       _UpperCamelCase   (			self )		->							Any:
			return (
			    MobileViTImageProcessor.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256' )
			    if is_vision_available()
			    else None
			)




	@slow
	def       _UpperCamelCase   (			self )		->							Dict:
			lowerCamelCase					:      Union[str, Any]          =   MobileViTVaForImageClassification.from_pretrained('apple/mobilevitv2-1.0-imagenet1k-256' ).to(
			    UpperCAmelCase_ )

			lowerCamelCase					:      Union[str, Any]          =   self.default_image_processor
			lowerCamelCase					:      Any          =   prepare_img()
			lowerCamelCase					:      str          =   image_processor(images=UpperCAmelCase_		,			return_tensors='pt' ).to(UpperCAmelCase_ )

			# forward pass
			with torch.no_grad():
					lowerCamelCase					:      Dict          =   model(**UpperCAmelCase_ )

			# verify the logits
			lowerCamelCase					:      List[Any]          =   torch.Size((1, 1000) )
			self.assertEqual(outputs.logits.shape		,			UpperCAmelCase_ )

			lowerCamelCase					:      Optional[int]          =   torch.tensor([-1.6_3_3_6E0_0, -7.3_2_0_4E-0_2, -5.1_8_8_3E-0_1] ).to(UpperCAmelCase_ )

			self.assertTrue(torch.allclose(outputs.logits[0, :3]		,			UpperCAmelCase_		,			atol=1E-4 ) )




	@slow
	def       _UpperCamelCase   (			self )		->							int:
			lowerCamelCase					:      Optional[Any]          =   MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )
			lowerCamelCase					:      Union[str, Any]          =   model.to(UpperCAmelCase_ )

			lowerCamelCase					:      List[str]          =   MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )

			lowerCamelCase					:      Any          =   prepare_img()
			lowerCamelCase					:      Tuple          =   image_processor(images=UpperCAmelCase_		,			return_tensors='pt' ).to(UpperCAmelCase_ )

			# forward pass
			with torch.no_grad():
					lowerCamelCase					:      Dict          =   model(**UpperCAmelCase_ )
			lowerCamelCase					:      Tuple          =   outputs.logits

			# verify the logits
			lowerCamelCase					:      str          =   torch.Size((1, 21, 32, 32) )
			self.assertEqual(logits.shape		,			UpperCAmelCase_ )

			lowerCamelCase					:      int          =   torch.tensor(
			    [
			        [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
			        [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
			        [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
			    ]		,			device=UpperCAmelCase_		,			)

			self.assertTrue(torch.allclose(logits[0, :3, :3, :3]		,			UpperCAmelCase_		,			atol=1E-4 ) )




	@slow
	def       _UpperCamelCase   (			self )		->							Any:
			lowerCamelCase					:      List[Any]          =   MobileViTVaForSemanticSegmentation.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )
			lowerCamelCase					:      Optional[int]          =   model.to(UpperCAmelCase_ )

			lowerCamelCase					:      Any          =   MobileViTImageProcessor.from_pretrained('shehan97/mobilevitv2-1.0-voc-deeplabv3' )

			lowerCamelCase					:      Tuple          =   prepare_img()
			lowerCamelCase					:      Any          =   image_processor(images=UpperCAmelCase_		,			return_tensors='pt' ).to(UpperCAmelCase_ )

			# forward pass
			with torch.no_grad():
					lowerCamelCase					:      Union[str, Any]          =   model(**UpperCAmelCase_ )

			lowerCamelCase					:      Optional[int]          =   outputs.logits.detach().cpu()

			lowerCamelCase					:      Optional[int]          =   image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_		,			target_sizes=[(50, 60)] )
			lowerCamelCase					:      List[str]          =   torch.Size((50, 60) )
			self.assertEqual(segmentation[0].shape		,			UpperCAmelCase_ )

			lowerCamelCase					:      Tuple          =   image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_ )
			lowerCamelCase					:      Optional[Any]          =   torch.Size((32, 32) )
			self.assertEqual(segmentation[0].shape		,			UpperCAmelCase_ )








"""simple docstring"""

import logging
import random

import ray

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


_A						=			logging.getLogger(__name__)






class 	_lowercase       :




	def __init__(			self )		->							Tuple:
			lowerCamelCase					:      Dict          =   False




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Any:
			if not self.initialized:
					lowerCamelCase					:      Optional[int]          =   RagRetriever(
					    UpperCAmelCase_		,			question_encoder_tokenizer=UpperCAmelCase_		,			generator_tokenizer=UpperCAmelCase_		,			index=UpperCAmelCase_		,			init_retrieval=UpperCAmelCase_		,			)
					lowerCamelCase					:      Union[str, Any]          =   True




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




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							str:
			lowerCamelCase     ,   lowerCamelCase					:      int          =   self.retriever._main_retrieve(UpperCAmelCase_		,			UpperCAmelCase_ )
			return doc_ids, retrieved_doc_embeds






class 	_lowercase       (    __UpperCAmelCase		):




	def __init__(			self		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_=None )		->							Dict:
			if index is not None and index.is_initialized() and len(UpperCAmelCase_ ) > 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__(
			    UpperCAmelCase_		,			question_encoder_tokenizer=UpperCAmelCase_		,			generator_tokenizer=UpperCAmelCase_		,			index=UpperCAmelCase_		,			init_retrieval=UpperCAmelCase_		,			)
			lowerCamelCase					:      List[str]          =   retrieval_workers
			if len(self.retrieval_workers ) > 0:
					ray.get(
					    [
					        worker.create_rag_retriever.remote(UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_ )
					        for worker in self.retrieval_workers
					    ] )




	def       _UpperCamelCase   (			self )		->							Dict:
			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       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							List[Any]:
			if len(self.retrieval_workers ) > 0:
					# Select a random retrieval actor.
					lowerCamelCase					:      Tuple          =   self.retrieval_workers[random.randint(0		,			len(self.retrieval_workers ) - 1 )]
					lowerCamelCase     ,   lowerCamelCase					:      Tuple          =   ray.get(random_worker.retrieve.remote(UpperCAmelCase_		,			UpperCAmelCase_ ) )
			else:
					lowerCamelCase     ,   lowerCamelCase					:      int          =   self._main_retrieve(UpperCAmelCase_		,			UpperCAmelCase_ )
			return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(UpperCAmelCase_ )




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




	@classmethod
	def       _UpperCamelCase   (			cls		,			UpperCAmelCase_		,			UpperCAmelCase_		,			UpperCAmelCase_=None		,			**UpperCAmelCase_ )		->							Dict:
			lowerCamelCase					:      str          =   kwargs.pop('config'		,			UpperCAmelCase_ ) or RagConfig.from_pretrained(UpperCAmelCase_		,			**UpperCAmelCase_ )
			lowerCamelCase					:      Any          =   RagTokenizer.from_pretrained(UpperCAmelCase_		,			config=UpperCAmelCase_ )
			lowerCamelCase					:      Any          =   rag_tokenizer.question_encoder
			lowerCamelCase					:      Union[str, Any]          =   rag_tokenizer.generator
			if indexed_dataset is not None:
					lowerCamelCase					:      int          =   'custom'
					lowerCamelCase					:      int          =   CustomHFIndex(config.retrieval_vector_size		,			UpperCAmelCase_ )
			else:
					lowerCamelCase					:      Optional[Any]          =   cls._build_index(UpperCAmelCase_ )
			return cls(
			    UpperCAmelCase_		,			question_encoder_tokenizer=UpperCAmelCase_		,			generator_tokenizer=UpperCAmelCase_		,			retrieval_workers=UpperCAmelCase_		,			index=UpperCAmelCase_		,			)






