
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        _UpperCamelCase				(					__UpperCAmelCase  ):






			'''simple docstring'''

			def       lowerCAmelCase__	(       self :       List[str]      ):
							UpperCamelCase_:							Tuple	=						self.config_class(**self.inputs_dict      )
							self.parent.assertTrue(hasattr(UpperCAmelCase_						,	"""width_multiplier"""      )      )




class        _UpperCamelCase				:






			'''simple docstring'''

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

			def       lowerCAmelCase__	(       self :       List[Any]      ):
							UpperCamelCase_:							int	=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

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

							UpperCamelCase_:							Union[str, Any]	=						self.get_config()

							return config, pixel_values, labels, pixel_labels

			def       lowerCAmelCase__	(       self :       Tuple      ):
							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       lowerCAmelCase__	(       self :       Optional[Any]						,	snake_case_ :       List[str]						,	snake_case_ :       Tuple						,	snake_case_ :       Dict						,	snake_case_ :       Union[str, Any]      ):
							UpperCamelCase_:							Union[str, Any]	=						MobileViTVaModel(config=UpperCAmelCase_      )
							model.to(UpperCAmelCase_      )
							model.eval()
							UpperCamelCase_:							Tuple	=						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       lowerCAmelCase__	(       self :       str						,	snake_case_ :       Union[str, Any]						,	snake_case_ :       Optional[Any]						,	snake_case_ :       Union[str, Any]						,	snake_case_ :       List[str]      ):
							UpperCamelCase_:							List[str]	=						self.num_labels
							UpperCamelCase_:							Any	=						MobileViTVaForImageClassification(UpperCAmelCase_      )
							model.to(UpperCAmelCase_      )
							model.eval()
							UpperCamelCase_:							List[str]	=						model(UpperCAmelCase_						,	labels=UpperCAmelCase_      )
							self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)      )

			def       lowerCAmelCase__	(       self :       Union[str, Any]						,	snake_case_ :       List[str]						,	snake_case_ :       Tuple						,	snake_case_ :       Optional[int]						,	snake_case_ :       Tuple      ):
							UpperCamelCase_:							Optional[int]	=						self.num_labels
							UpperCamelCase_:							int	=						MobileViTVaForSemanticSegmentation(UpperCAmelCase_      )
							model.to(UpperCAmelCase_      )
							model.eval()
							UpperCamelCase_:							int	=						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,
							    )						,	)
							UpperCamelCase_:							Optional[Any]	=						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       lowerCAmelCase__	(       self :       List[Any]      ):
							UpperCamelCase_:							Optional[Any]	=						self.prepare_config_and_inputs()
							UpperCamelCase_,   UpperCamelCase_,   UpperCamelCase_,   UpperCamelCase_:							Optional[Any]	=						config_and_inputs
							UpperCamelCase_:							int	=						{"""pixel_values""": pixel_values}
							return config, inputs_dict




@require_torch
class        _UpperCamelCase				(					__UpperCAmelCase     ,       __UpperCAmelCase     ,       unittest.TestCase  ):






			'''simple docstring'''

			__UpperCamelCase      :						List[str]					      = (
			    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
			    if is_torch_available()
			    else ()
			)

			__UpperCamelCase      :						Union[str, Any]					      = (
			    {
			        """feature-extraction""": MobileViTVaModel,
			        """image-classification""": MobileViTVaForImageClassification,
			        """image-segmentation""": MobileViTVaForSemanticSegmentation,
			    }
			    if is_torch_available()
			    else {}
			)

			__UpperCamelCase      :						str					      = False
			__UpperCamelCase      :						Dict					      = False
			__UpperCamelCase      :						Union[str, Any]					      = False
			__UpperCamelCase      :						Optional[int]					      = False

			def       lowerCAmelCase__	(       self :       List[str]      ):
							UpperCamelCase_:							Union[str, Any]	=						MobileViTVaModelTester(self      )
							UpperCamelCase_:							Any	=						MobileViTVaConfigTester(self						,	config_class=UpperCAmelCase_						,	has_text_modality=UpperCAmelCase_      )

			def       lowerCAmelCase__	(       self :       Optional[Any]      ):
							self.config_tester.run_common_tests()

			@unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""      )
			def       lowerCAmelCase__	(       self :       str      ):
							pass

			@unittest.skip(reason="""MobileViTV2 does not support input and output embeddings"""      )
			def       lowerCAmelCase__	(       self :       Union[str, Any]      ):
							pass

			@unittest.skip(reason="""MobileViTV2 does not output attentions"""      )
			def       lowerCAmelCase__	(       self :       Tuple      ):
							pass

			@require_torch_multi_gpu
			@unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""      )
			def       lowerCAmelCase__	(       self :       Tuple      ):
							pass

			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""      )
			def       lowerCAmelCase__	(       self :       Any      ):
							pass

			def       lowerCAmelCase__	(       self :       Any      ):
							UpperCamelCase_,   UpperCamelCase_:							Optional[int]	=						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											UpperCamelCase_:							Optional[Any]	=						model_class(UpperCAmelCase_      )
											UpperCamelCase_:							Any	=						inspect.signature(model.forward      )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											UpperCamelCase_:							List[str]	=						[*signature.parameters.keys()]

											UpperCamelCase_:							Optional[int]	=						["""pixel_values"""]
											self.assertListEqual(arg_names[:1]						,	UpperCAmelCase_      )

			def       lowerCAmelCase__	(       self :       Optional[int]      ):
							UpperCamelCase_:							List[str]	=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase_      )

			def       lowerCAmelCase__	(       self :       Optional[Any]      ):
							def check_hidden_states_output(snake_case_ :       int						,	snake_case_ :       Dict						,	snake_case_ :       Any      ):
											UpperCamelCase_:							Dict	=						model_class(UpperCAmelCase_      )
											model.to(UpperCAmelCase_      )
											model.eval()

											with torch.no_grad():
															UpperCamelCase_:							Dict	=						model(**self._prepare_for_class(UpperCAmelCase_						,	UpperCAmelCase_      )      )

											UpperCamelCase_:							List[str]	=						outputs.hidden_states

											UpperCamelCase_:							Any	=						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.
											UpperCamelCase_:							Dict	=						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      )

							UpperCamelCase_,   UpperCamelCase_:							List[str]	=						self.model_tester.prepare_config_and_inputs_for_common()

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

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											UpperCamelCase_:							Optional[int]	=						True

											check_hidden_states_output(UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_      )

			def       lowerCAmelCase__	(       self :       List[str]      ):
							UpperCamelCase_:							Tuple	=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_      )

			def       lowerCAmelCase__	(       self :       str      ):
							UpperCamelCase_:							str	=						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*UpperCAmelCase_      )




			@slow
			def       lowerCAmelCase__	(       self :       List[Any]      ):
							for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											UpperCamelCase_:							Dict	=						MobileViTVaModel.from_pretrained(UpperCAmelCase_      )
											self.assertIsNotNone(UpperCAmelCase_      )







def        A__					(						)  ->	Tuple:

				'''simple docstring'''





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





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






			'''simple docstring'''

			@cached_property
			def       lowerCAmelCase__	(       self :       Tuple      ):
							return (
							    MobileViTImageProcessor.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""      )
							    if is_vision_available()
							    else None
							)

			@slow
			def       lowerCAmelCase__	(       self :       Any      ):
							UpperCamelCase_:							int	=						MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""      ).to(
							    UpperCAmelCase_      )

							UpperCamelCase_:							Any	=						self.default_image_processor
							UpperCamelCase_:							str	=						prepare_img()
							UpperCamelCase_:							Tuple	=						image_processor(images=UpperCAmelCase_						,	return_tensors="""pt"""      ).to(UpperCAmelCase_      )

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

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

							UpperCamelCase_:							str	=						torch.tensor([-1.6_336e00, -7.3_204e-02, -5.1_883e-01]      ).to(UpperCAmelCase_      )

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

			@slow
			def       lowerCAmelCase__	(       self :       Any      ):
							UpperCamelCase_:							int	=						MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""      )
							UpperCamelCase_:							Tuple	=						model.to(UpperCAmelCase_      )

							UpperCamelCase_:							Optional[Any]	=						MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""      )

							UpperCamelCase_:							str	=						prepare_img()
							UpperCamelCase_:							Optional[int]	=						image_processor(images=UpperCAmelCase_						,	return_tensors="""pt"""      ).to(UpperCAmelCase_      )

							# forward pass
							with torch.no_grad():
											UpperCamelCase_:							Any	=						model(**UpperCAmelCase_      )
							UpperCamelCase_:							Any	=						outputs.logits

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

							UpperCamelCase_:							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       lowerCAmelCase__	(       self :       Tuple      ):
							UpperCamelCase_:							Tuple	=						MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""      )
							UpperCamelCase_:							Tuple	=						model.to(UpperCAmelCase_      )

							UpperCamelCase_:							str	=						MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""      )

							UpperCamelCase_:							Optional[Any]	=						prepare_img()
							UpperCamelCase_:							List[Any]	=						image_processor(images=UpperCAmelCase_						,	return_tensors="""pt"""      ).to(UpperCAmelCase_      )

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

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

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

							UpperCamelCase_:							List[Any]	=						image_processor.post_process_semantic_segmentation(outputs=UpperCAmelCase_      )
							UpperCamelCase_:							str	=						torch.Size((32, 32)      )
							self.assertEqual(segmentation[0].shape						,	UpperCAmelCase_      )








import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
lowerCamelCase_						: Union[str, Any]         =  logging.getLogger()

lowerCamelCase_						: List[str]         =  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)

class        _UpperCamelCase				(					_A  ):






			'''simple docstring'''

			def       lowerCAmelCase__	(       self :       List[Any]						,	snake_case_ :       Dict      ):
							os.makedirs(snake_case_						,	exist_ok=snake_case_      )
							UpperCamelCase_:							int	=						{"""source""": """What is love ?""", """target""": """life"""}
							UpperCamelCase_:							Tuple	=						{"""train""": 12, """val""": 2, """test""": 2}
							for split in ["train", "test", "val"]:
											for field in ["source", "target"]:
															UpperCamelCase_:							Tuple	=						"""\n""".join([contents[field]] * n_lines[split]      )
															with open(os.path.join(snake_case_						,	f'''{split}.{field}'''      )						,	"""w"""      ) as f:
																			f.write(snake_case_      )

			def       lowerCAmelCase__	(       self :       Dict						,	snake_case_ :       int						,	snake_case_ :       str = "pytorch"      ):
							UpperCamelCase_:							Optional[Any]	=						self.get_auto_remove_tmp_dir()
							UpperCamelCase_:							Dict	=						os.path.join(snake_case_						,	"""output"""      )
							UpperCamelCase_:							Any	=						os.path.join(snake_case_						,	"""data"""      )
							self._create_dummy_data(data_dir=snake_case_      )

							UpperCamelCase_:							Union[str, Any]	=						f'''
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            '''.split()

							if gpus > 0:
											testargs.append(f'''--gpus={gpus}'''      )
											if is_apex_available():
															testargs.append("""--fp16"""      )
							else:
											testargs.append("""--gpus=0"""      )
											testargs.append("""--distributed_backend=ddp_cpu"""      )
											testargs.append("""--num_processes=2"""      )

							UpperCamelCase_:							Optional[Any]	=						[sys.executable, str(Path(finetune_rag.__file__      ).resolve()      )] + testargs
							execute_subprocess_async(snake_case_						,	env=self.get_env()      )

							UpperCamelCase_:							Optional[int]	=						os.path.join(snake_case_						,	"""metrics.json"""      )
							with open(snake_case_      ) as f:
											UpperCamelCase_:							Any	=						json.load(snake_case_      )
							return result

			@require_torch_gpu
			def       lowerCAmelCase__	(       self :       Optional[Any]      ):
							UpperCamelCase_:							List[str]	=						self._run_finetune(gpus=1      )
							self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]						,	0.2      )

			@require_torch_multi_gpu
			def       lowerCAmelCase__	(       self :       Optional[Any]      ):
							UpperCamelCase_:							Any	=						self._run_finetune(gpus=2      )
							self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]						,	0.2      )

			@require_torch_gpu
			@require_ray
			def       lowerCAmelCase__	(       self :       Dict      ):
							UpperCamelCase_:							List[Any]	=						self._run_finetune(gpus=1						,	distributed_retriever="""ray"""      )
							self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]						,	0.2      )




			@require_torch_multi_gpu
			@require_ray
			def       lowerCAmelCase__	(       self :       Dict      ):
							UpperCamelCase_:							List[Any]	=						self._run_finetune(gpus=1						,	distributed_retriever="""ray"""      )
							self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]						,	0.2      )






