
"""simple docstring"""






import argparse
from typing import List

import evaluate
import numpy as np
import torch
from datasets import DatasetDict, load_dataset

# New Code #
# We'll be using StratifiedKFold for this example
from sklearn.model_selection import StratifiedKFold
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to perform Cross Validation,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowercase							=				16
_lowercase							=				32



def      lowerCAmelCase__			(		__magic_name__			,   __magic_name__			,   __magic_name__			,   __magic_name__			,   __magic_name__ = 1_6      )						->Any:
			__lowercase       =       AutoTokenizer.from_pretrained("bert-base-cased"      )
			__lowercase       =       DatasetDict(
			    {
			        "train": dataset["train"].select(__magic_name__      ),
			        "validation": dataset["train"].select(__magic_name__      ),
			        "test": dataset["validation"],
			    }      )

			def tokenize_function(__magic_name__      ):
						# max_length=None => use the model max length (it's actually the default)
						__lowercase       =       tokenizer(examples["sentence1"]			,   examples["sentence2"]			,   truncation=__magic_name__			,   max_length=__magic_name__      )
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						__lowercase       =       datasets.map(
						    __magic_name__			,   batched=__magic_name__			,   remove_columns=["idx", "sentence1", "sentence2"]			,   )

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			__lowercase       =       tokenized_datasets.rename_column("label"			,   "labels"      )

			def collate_fn(__magic_name__      ):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						__lowercase       =       1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									__lowercase       =       1_6
						elif accelerator.mixed_precision != "no":
									__lowercase       =       8
						else:
									__lowercase       =       None

						return tokenizer.pad(
						    __magic_name__			,   padding="longest"			,   max_length=__magic_name__			,   pad_to_multiple_of=__magic_name__			,   return_tensors="pt"			,   )

			# Instantiate dataloaders.
			__lowercase       =       DataLoader(
			    tokenized_datasets["train"]			,   shuffle=__magic_name__			,   collate_fn=__magic_name__			,   batch_size=__magic_name__      )
			__lowercase       =       DataLoader(
			    tokenized_datasets["validation"]			,   shuffle=__magic_name__			,   collate_fn=__magic_name__			,   batch_size=__magic_name__      )

			__lowercase       =       DataLoader(
			    tokenized_datasets["test"]			,   shuffle=__magic_name__			,   collate_fn=__magic_name__			,   batch_size=__magic_name__      )

			return train_dataloader, eval_dataloader, test_dataloader



def      lowerCAmelCase__			(		__magic_name__			,   __magic_name__      )						->int:
			# New Code #
			__lowercase       =       []
			# Download the dataset
			__lowercase       =       load_dataset("glue"			,   "mrpc"      )
			# Create our splits
			__lowercase       =       StratifiedKFold(n_splits=int(args.num_folds      )      )
			# Initialize accelerator
			__lowercase       =       Accelerator(cpu=args.cpu			,   mixed_precision=args.mixed_precision      )
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			__lowercase       =       config["lr"]
			__lowercase       =       int(config["num_epochs"]      )
			__lowercase       =       int(config["seed"]      )
			__lowercase       =       int(config["batch_size"]      )

			__lowercase       =       evaluate.load("glue"			,   "mrpc"      )

			# If the batch size is too big we use gradient accumulation
			__lowercase       =       1
			if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
						__lowercase       =       batch_size // MAX_GPU_BATCH_SIZE
						__lowercase       =       MAX_GPU_BATCH_SIZE

			set_seed(__magic_name__      )

			# New Code #
			# Create our folds:
			__lowercase       =       kfold.split(np.zeros(datasets["train"].num_rows      )			,   datasets["train"]["label"]      )
			__lowercase       =       []
			# Iterate over them
			for i, (train_idxs, valid_idxs) in enumerate(__magic_name__      ):
						__lowercase   ,					__lowercase   ,					__lowercase       =       get_fold_dataloaders(
						    __magic_name__			,   __magic_name__			,   __magic_name__			,   __magic_name__			,   )
						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						__lowercase       =       AutoModelForSequenceClassification.from_pretrained("bert-base-cased"			,   return_dict=__magic_name__      )

						# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
						# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
						# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
						__lowercase       =       model.to(accelerator.device      )

						# Instantiate optimizer
						__lowercase       =       AdamW(params=model.parameters()			,   lr=__magic_name__      )

						# Instantiate scheduler
						__lowercase       =       get_linear_schedule_with_warmup(
						    optimizer=__magic_name__			,   num_warmup_steps=1_0_0			,   num_training_steps=(len(__magic_name__      ) * num_epochs) // gradient_accumulation_steps			,   )

						# Prepare everything
						# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
						# prepare method.
						__lowercase   ,					__lowercase   ,					__lowercase   ,					__lowercase   ,					__lowercase       =       accelerator.prepare(
						    __magic_name__			,   __magic_name__			,   __magic_name__			,   __magic_name__			,   __magic_name__      )

						# Now we train the model
						for epoch in range(__magic_name__      ):
									model.train()
									for step, batch in enumerate(__magic_name__      ):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device      )
												__lowercase       =       model(**__magic_name__      )
												__lowercase       =       outputs.loss
												__lowercase       =       loss / gradient_accumulation_steps
												accelerator.backward(__magic_name__      )
												if step % gradient_accumulation_steps == 0:
															optimizer.step()
															lr_scheduler.step()
															optimizer.zero_grad()

									model.eval()
									for step, batch in enumerate(__magic_name__      ):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device      )
												with torch.no_grad():
															__lowercase       =       model(**__magic_name__      )
												__lowercase       =       outputs.logits.argmax(dim=-1      )
												__lowercase   ,					__lowercase       =       accelerator.gather_for_metrics((predictions, batch["labels"])      )
												metric.add_batch(
												    predictions=__magic_name__			,   references=__magic_name__			,   )

									__lowercase       =       metric.compute()
									# Use accelerator.print to print only on the main process.
									accelerator.print(F'''epoch {epoch}:'''			,   __magic_name__      )

						# New Code #
						# We also run predictions on the test set at the very end
						__lowercase       =       []
						for step, batch in enumerate(__magic_name__      ):
									# We could avoid this line since we set the accelerator with `device_placement=True`.
									batch.to(accelerator.device      )
									with torch.no_grad():
												__lowercase       =       model(**__magic_name__      )
									__lowercase       =       outputs.logits
									__lowercase   ,					__lowercase       =       accelerator.gather_for_metrics((predictions, batch["labels"])      )
									fold_predictions.append(predictions.cpu()      )
									if i == 0:
												# We need all of the test predictions
												test_references.append(references.cpu()      )
        # Use accelerator.print to print only on the main process.
						test_predictions.append(torch.cat(__magic_name__			,   dim=0      )      )
						# We now need to release all our memory and get rid of the current model, optimizer, etc
						accelerator.free_memory()
			# New Code #
			# Finally we check the accuracy of our folded results:
			__lowercase       =       torch.cat(__magic_name__			,   dim=0      )
			__lowercase       =       torch.stack(__magic_name__			,   dim=0      ).sum(dim=0      ).div(int(args.num_folds      )      ).argmax(dim=-1      )
			__lowercase       =       metric.compute(predictions=__magic_name__			,   references=__magic_name__      )
			accelerator.print("Average test metrics from all folds:"			,   __magic_name__      )



def      lowerCAmelCase__			(		)						->Union[str, Any]:
			__lowercase       =       argparse.ArgumentParser(description="Simple example of training script."      )
			parser.add_argument(
			    "--mixed_precision"			,   type=__magic_name__			,   default=__magic_name__			,   choices=["no", "fp16", "bf16", "fp8"]			,   help="Whether to use mixed precision. Choose"
			    "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
			    "and an Nvidia Ampere GPU."			,   )
			parser.add_argument("--cpu"			,   action="store_true"			,   help="If passed, will train on the CPU."      )
			# New Code #
			parser.add_argument("--num_folds"			,   type=__magic_name__			,   default=3			,   help="The number of splits to perform across the dataset"      )
			__lowercase       =       parser.parse_args()
			__lowercase       =       {"lr": 2E-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
			training_function(__magic_name__			,   __magic_name__      )


if __name__ == "__main__":
					main()







"""simple docstring"""






import inspect
import unittest

from transformers import MobileViTConfig
from transformers.testing_utils import require_torch, 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 MobileViTForImageClassification, MobileViTForSemanticSegmentation, MobileViTModel
					from transformers.models.mobilevit.modeling_mobilevit import MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image

					from transformers import MobileViTImageProcessor

class 					__a     ( __a       ):



							'''simple docstring'''







							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[str]:




										'''simple docstring'''
										__lowercase       =       self.config_class(**self.inputs_dict       )
										self.parent.assertTrue(hasattr(_lowerCamelCase  ,    "hidden_sizes"       )       )
										self.parent.assertTrue(hasattr(_lowerCamelCase  ,    "neck_hidden_sizes"       )       )
										self.parent.assertTrue(hasattr(_lowerCamelCase  ,    "num_attention_heads"       )       )





class 					__a     :



							'''simple docstring'''







							def __init__(     self  ,    _lowerCamelCase  ,    _lowerCamelCase=13  ,    _lowerCamelCase=32  ,    _lowerCamelCase=2  ,    _lowerCamelCase=3  ,    _lowerCamelCase=640  ,    _lowerCamelCase=4  ,    _lowerCamelCase="silu"  ,    _lowerCamelCase=3  ,    _lowerCamelCase=32  ,    _lowerCamelCase=0.1  ,    _lowerCamelCase=0.1  ,    _lowerCamelCase=0.1  ,    _lowerCamelCase=0.02  ,    _lowerCamelCase=True  ,    _lowerCamelCase=True  ,    _lowerCamelCase=10  ,    _lowerCamelCase=None  ,    )			->							Tuple:




										'''simple docstring'''
										__lowercase       =       parent
										__lowercase       =       batch_size
										__lowercase       =       image_size
										__lowercase       =       patch_size
										__lowercase       =       num_channels
										__lowercase       =       last_hidden_size
										__lowercase       =       num_attention_heads
										__lowercase       =       hidden_act
										__lowercase       =       conv_kernel_size
										__lowercase       =       output_stride
										__lowercase       =       hidden_dropout_prob
										__lowercase       =       attention_probs_dropout_prob
										__lowercase       =       classifier_dropout_prob
										__lowercase       =       use_labels
										__lowercase       =       is_training
										__lowercase       =       num_labels
										__lowercase       =       initializer_range
										__lowercase       =       scope
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Tuple:




										'''simple docstring'''
										__lowercase       =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]       )

										__lowercase       =       None
										__lowercase       =       None
										if self.use_labels:
													__lowercase       =       ids_tensor([self.batch_size]  ,    self.num_labels       )
													__lowercase       =       ids_tensor([self.batch_size, self.image_size, self.image_size]  ,    self.num_labels       )

										__lowercase       =       self.get_config()

										return config, pixel_values, labels, pixel_labels
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Tuple:




										'''simple docstring'''
										return MobileViTConfig(
										    image_size=self.image_size  ,    patch_size=self.patch_size  ,    num_channels=self.num_channels  ,    num_attention_heads=self.num_attention_heads  ,    hidden_act=self.hidden_act  ,    conv_kernel_size=self.conv_kernel_size  ,    output_stride=self.output_stride  ,    hidden_dropout_prob=self.hidden_dropout_prob  ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,    classifier_dropout_prob=self.classifier_dropout_prob  ,    initializer_range=self.initializer_range  ,    )
							def   SCREAMING_SNAKE_CASE						(     self  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )			->							str:




										'''simple docstring'''
										__lowercase       =       MobileViTModel(config=_lowerCamelCase       )
										model.to(_lowerCamelCase       )
										model.eval()
										__lowercase       =       model(_lowerCamelCase       )
										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   SCREAMING_SNAKE_CASE						(     self  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )			->							int:




										'''simple docstring'''
										__lowercase       =       self.num_labels
										__lowercase       =       MobileViTForImageClassification(_lowerCamelCase       )
										model.to(_lowerCamelCase       )
										model.eval()
										__lowercase       =       model(_lowerCamelCase  ,    labels=_lowerCamelCase       )
										self.parent.assertEqual(result.logits.shape  ,    (self.batch_size, self.num_labels)       )
							def   SCREAMING_SNAKE_CASE						(     self  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )			->							Optional[int]:




										'''simple docstring'''
										__lowercase       =       self.num_labels
										__lowercase       =       MobileViTForSemanticSegmentation(_lowerCamelCase       )
										model.to(_lowerCamelCase       )
										model.eval()
										__lowercase       =       model(_lowerCamelCase       )
										self.parent.assertEqual(
										    result.logits.shape  ,    (
										        self.batch_size,
										        self.num_labels,
										        self.image_size // self.output_stride,
										        self.image_size // self.output_stride,
										    )  ,    )
										__lowercase       =       model(_lowerCamelCase  ,    labels=_lowerCamelCase       )
										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   SCREAMING_SNAKE_CASE						(     self       )			->							List[str]:




										'''simple docstring'''
										__lowercase       =       self.prepare_config_and_inputs()
										__lowercase   ,					__lowercase   ,					__lowercase   ,					__lowercase       =       config_and_inputs
										__lowercase       =       {"pixel_values": pixel_values}
										return config, inputs_dict





@require_torch
class 					__a     ( __a					,   __a					,   unittest.TestCase       ):



							'''simple docstring'''







							_lowerCamelCase					:			str          =   (
							    (MobileViTModel, MobileViTForImageClassification, MobileViTForSemanticSegmentation)
							    if is_torch_available()
							    else ()
							)
							_lowerCamelCase					:			Optional[int]          =   (
							    {
							        """feature-extraction""": MobileViTModel,
							        """image-classification""": MobileViTForImageClassification,
							        """image-segmentation""": MobileViTForSemanticSegmentation,
							    }
							    if is_torch_available()
							    else {}
							)

							_lowerCamelCase					:			Tuple          =   False
							_lowerCamelCase					:			Optional[Any]          =   False
							_lowerCamelCase					:			Any          =   False
							_lowerCamelCase					:			List[Any]          =   False
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Tuple:




										'''simple docstring'''
										__lowercase       =       MobileViTModelTester(self       )
										__lowercase       =       MobileViTConfigTester(self  ,    config_class=_lowerCamelCase  ,    has_text_modality=_lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Dict:




										'''simple docstring'''
										self.config_tester.run_common_tests()
							@unittest.skip(reason="MobileViT does not use inputs_embeds"       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							str:




										'''simple docstring'''
										pass
							@unittest.skip(reason="MobileViT does not support input and output embeddings"       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Dict:




										'''simple docstring'''
										pass
							@unittest.skip(reason="MobileViT does not output attentions"       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Dict:




										'''simple docstring'''
										pass
							def   SCREAMING_SNAKE_CASE						(     self       )			->							int:




										'''simple docstring'''
										__lowercase   ,					__lowercase       =       self.model_tester.prepare_config_and_inputs_for_common()

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

													__lowercase       =       ["pixel_values"]
													self.assertListEqual(arg_names[:1]  ,    _lowerCamelCase       )
							@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Dict:




										'''simple docstring'''
										pass
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Tuple:




										'''simple docstring'''
										__lowercase       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[Any]:




										'''simple docstring'''
										def check_hidden_states_output(_lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       ):
													__lowercase       =       model_class(_lowerCamelCase       )
													model.to(_lowerCamelCase       )
													model.eval()

													with torch.no_grad():
																__lowercase       =       model(**self._prepare_for_class(_lowerCamelCase  ,    _lowerCamelCase       )       )

													__lowercase       =       outputs.hidden_states

													__lowercase       =       5
													self.assertEqual(len(_lowerCamelCase       )  ,    _lowerCamelCase       )

													# MobileViT's feature maps are of shape (batch_size, num_channels, height, width)
													# with the width and height being successively divided by 2.
													__lowercase       =       2
													for i in range(len(_lowerCamelCase       )       ):
																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       )

										__lowercase   ,					__lowercase       =       self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
													__lowercase       =       True
													check_hidden_states_output(_lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													__lowercase       =       True

													check_hidden_states_output(_lowerCamelCase  ,    _lowerCamelCase  ,    _lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[str]:




										'''simple docstring'''
										__lowercase       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase       )
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[Any]:




										'''simple docstring'''
										__lowercase       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_semantic_segmentation(*_lowerCamelCase       )





							@slow
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Tuple:




										'''simple docstring'''
										for model_name in MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__lowercase       =       MobileViTModel.from_pretrained(_lowerCamelCase       )
													self.assertIsNotNone(_lowerCamelCase       )





def      lowerCAmelCase__			(		)						->List[Any]:
			__lowercase       =       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"      )
			return image







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



							'''simple docstring'''







							@cached_property
							def   SCREAMING_SNAKE_CASE						(     self       )			->							str:




										'''simple docstring'''
										return MobileViTImageProcessor.from_pretrained("apple/mobilevit-xx-small"       ) if is_vision_available() else None
							@slow
							def   SCREAMING_SNAKE_CASE						(     self       )			->							Optional[int]:




										'''simple docstring'''
										__lowercase       =       MobileViTForImageClassification.from_pretrained("apple/mobilevit-xx-small"       ).to(_lowerCamelCase       )

										__lowercase       =       self.default_image_processor
										__lowercase       =       prepare_img()
										__lowercase       =       image_processor(images=_lowerCamelCase  ,    return_tensors="pt"       ).to(_lowerCamelCase       )

										# forward pass
										with torch.no_grad():
													__lowercase       =       model(**_lowerCamelCase       )

										# verify the logits
										__lowercase       =       torch.Size((1, 1_000)       )
										self.assertEqual(outputs.logits.shape  ,    _lowerCamelCase       )

										__lowercase       =       torch.tensor([-1.9364, -1.2327, -0.4653]       ).to(_lowerCamelCase       )

										self.assertTrue(torch.allclose(outputs.logits[0, :3]  ,    _lowerCamelCase  ,    atol=1e-4       )       )
							@slow
							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[str]:




										'''simple docstring'''
										__lowercase       =       MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small"       )
										__lowercase       =       model.to(_lowerCamelCase       )

										__lowercase       =       MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small"       )

										__lowercase       =       prepare_img()
										__lowercase       =       image_processor(images=_lowerCamelCase  ,    return_tensors="pt"       ).to(_lowerCamelCase       )

										# forward pass
										with torch.no_grad():
													__lowercase       =       model(**_lowerCamelCase       )
										__lowercase       =       outputs.logits

										# verify the logits
										__lowercase       =       torch.Size((1, 21, 32, 32)       )
										self.assertEqual(logits.shape  ,    _lowerCamelCase       )

										__lowercase       =       torch.tensor(
										    [
										        [[6.9713, 6.9786, 7.2422], [7.2893, 7.2825, 7.4446], [7.6580, 7.8797, 7.9420]],
										        [[-10.6869, -10.3250, -10.3471], [-10.4228, -9.9868, -9.7132], [-11.0405, -11.0221, -10.7318]],
										        [[-3.3089, -2.8539, -2.6740], [-3.2706, -2.5621, -2.5108], [-3.2534, -2.6615, -2.6651]],
										    ]  ,    device=_lowerCamelCase  ,    )

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





							@slow
							def   SCREAMING_SNAKE_CASE						(     self       )			->							List[str]:




										'''simple docstring'''
										__lowercase       =       MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small"       )
										__lowercase       =       model.to(_lowerCamelCase       )

										__lowercase       =       MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small"       )

										__lowercase       =       prepare_img()
										__lowercase       =       image_processor(images=_lowerCamelCase  ,    return_tensors="pt"       ).to(_lowerCamelCase       )

										# forward pass
										with torch.no_grad():
													__lowercase       =       model(**_lowerCamelCase       )

										__lowercase       =       outputs.logits.detach().cpu()

										__lowercase       =       image_processor.post_process_semantic_segmentation(outputs=_lowerCamelCase  ,    target_sizes=[(50, 60)]       )
										__lowercase       =       torch.Size((50, 60)       )
										self.assertEqual(segmentation[0].shape  ,    _lowerCamelCase       )

										__lowercase       =       image_processor.post_process_semantic_segmentation(outputs=_lowerCamelCase       )
										__lowercase       =       torch.Size((32, 32)       )
										self.assertEqual(segmentation[0].shape  ,    _lowerCamelCase       )





