


"""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    lowerCamelCase__						(		_a    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       __lowercase									=	self.config_class(**self.inputs_dict							)
       self.parent.assertTrue(hasattr(A_      ,       """width_multiplier"""							)							)






class    lowerCamelCase__						:
      def __init__(		self	:	int      ,       A_	:	str      ,       A_	:	Tuple=1_3      ,       A_	:	Tuple=6_4      ,       A_	:	int=2      ,       A_	:	Optional[Any]=3      ,       A_	:	Union[str, Any]="swish"      ,       A_	:	Optional[int]=3      ,       A_	:	List[str]=3_2      ,       A_	:	Dict=0.1      ,       A_	:	List[Any]=0.02      ,       A_	:	Any=True      ,       A_	:	Tuple=True      ,       A_	:	Dict=1_0      ,       A_	:	Any=None      ,       A_	:	Any=0.25      ,       A_	:	List[Any]=0.0      ,       A_	:	Tuple=0.0      ,       ):






       '''simple docstring'''


       __lowercase									=	parent
       __lowercase									=	batch_size
       __lowercase									=	image_size
       __lowercase									=	patch_size
       __lowercase									=	num_channels
       __lowercase									=	make_divisible(5_1_2 * width_multiplier      ,       divisor=8							)
       __lowercase									=	hidden_act
       __lowercase									=	conv_kernel_size
       __lowercase									=	output_stride
       __lowercase									=	classifier_dropout_prob
       __lowercase									=	use_labels
       __lowercase									=	is_training
       __lowercase									=	num_labels
       __lowercase									=	initializer_range
       __lowercase									=	scope
       __lowercase									=	width_multiplier
       __lowercase									=	ffn_dropout
       __lowercase									=	attn_dropout
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''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	:	Union[str, Any]							):






       '''simple docstring'''


       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 SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]      ,       A_	:	int      ,       A_	:	Dict      ,       A_	:	Tuple      ,       A_	:	int							):






       '''simple docstring'''


       __lowercase									=	MobileViTVaModel(config=A_							)
       model.to(A_							)
       model.eval()
       __lowercase									=	model(A_							)
       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	:	str      ,       A_	:	str      ,       A_	:	Dict      ,       A_	:	int      ,       A_	:	Tuple							):






       '''simple docstring'''


       __lowercase									=	self.num_labels
       __lowercase									=	MobileViTVaForImageClassification(A_							)
       model.to(A_							)
       model.eval()
       __lowercase									=	model(A_      ,       labels=A_							)
       self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)							)
      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	Union[str, Any]      ,       A_	:	Any      ,       A_	:	Dict      ,       A_	:	Dict							):






       '''simple docstring'''


       __lowercase									=	self.num_labels
       __lowercase									=	MobileViTVaForSemanticSegmentation(A_							)
       model.to(A_							)
       model.eval()
       __lowercase									=	model(A_							)
       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(A_      ,       labels=A_							)
       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	:	Dict							):






       '''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    lowerCamelCase__						(		_a			,   _a			,   unittest.TestCase    ):
      a :     Optional[Any]											=		(
          (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
          if is_torch_available()
          else ()
      )

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

      a :     List[Any]											=		False
      a :     Union[str, Any]											=		False
      a :     Optional[Any]											=		False
      a :     Union[str, Any]											=		False
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	MobileViTVaModelTester(self							)
       __lowercase									=	MobileViTVaConfigTester(self      ,       config_class=A_      ,       has_text_modality=A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       self.config_tester.run_common_tests()
      @unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


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






       '''simple docstring'''


       pass
      @unittest.skip(reason="""MobileViTV2 does not output attentions"""							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


       pass
      @require_torch_multi_gpu
      @unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


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






       '''simple docstring'''


       pass
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       __lowercase							,		__lowercase									=	self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        __lowercase									=	model_class(A_							)
        __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]      ,       A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       __lowercase									=	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       def check_hidden_states_output(A_	:	str      ,       A_	:	Tuple      ,       A_	:	Union[str, Any]							):
        __lowercase									=	model_class(A_							)
        model.to(A_							)
        model.eval()

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

        __lowercase									=	outputs.hidden_states

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

        # MobileViTV2'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(A_							)							):
         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(A_      ,       A_      ,       A_							)

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

        check_hidden_states_output(A_      ,       A_      ,       A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


       __lowercase									=	self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*A_							)
      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(*A_							)






      @slow
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowercase									=	MobileViTVaModel.from_pretrained(A_							)
        self.assertIsNotNone(A_							)






def lowerCAmelCase_      (				):






 """simple docstring"""







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





@require_torch
@require_vision
class    lowerCamelCase__						(		unittest.TestCase    ):
      @cached_property
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       return (
           MobileViTImageProcessor.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""							)
           if is_vision_available()
           else None
       )
      @slow
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""							).to(
           A_							)

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

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

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

       __lowercase									=	torch.tensor([-1.6_3_3_6e0_0, -7.3_2_0_4e-0_2, -5.1_8_8_3e-0_1]							).to(A_							)

       self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,       A_      ,       atol=1e-4							)							)
      @slow
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)
       __lowercase									=	model.to(A_							)

       __lowercase									=	MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)

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

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

       # verify the logits
       __lowercase									=	torch.Size((1, 2_1, 3_2, 3_2)							)
       self.assertEqual(logits.shape      ,       A_							)

       __lowercase									=	torch.tensor(
           [
               [[7.08_63, 7.15_25, 6.82_01], [6.69_31, 6.87_70, 6.89_33], [6.29_78, 7.03_66, 6.96_36]],
               [[-3.71_34, -3.67_12, -3.66_75], [-3.58_25, -3.35_49, -3.47_77], [-3.34_35, -3.39_79, -3.28_57]],
               [[-2.93_29, -2.80_03, -2.73_69], [-3.05_64, -2.47_80, -2.02_07], [-2.68_89, -1.92_98, -1.76_40]],
           ]      ,       device=A_      ,       )

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






      @slow
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


       __lowercase									=	MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)
       __lowercase									=	model.to(A_							)

       __lowercase									=	MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""							)

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

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

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

       __lowercase									=	image_processor.post_process_semantic_segmentation(outputs=A_      ,       target_sizes=[(5_0, 6_0)]							)
       __lowercase									=	torch.Size((5_0, 6_0)							)
       self.assertEqual(segmentation[0].shape      ,       A_							)

       __lowercase									=	image_processor.post_process_semantic_segmentation(outputs=A_							)
       __lowercase									=	torch.Size((3_2, 3_2)							)
       self.assertEqual(segmentation[0].shape      ,       A_							)




"""simple docstring"""

import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO
)
UpperCAmelCase__  							=logging.getLogger(__name__)

if __name__ == "__main__":
       UpperCAmelCase__  							=argparse.ArgumentParser(
           description="Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)"
       )
       parser.add_argument(
           "--data_file", type=str, default="data/dump.bert-base-uncased.pickle", help="The binarized dataset."
       )
       parser.add_argument(
           "--token_counts_dump", type=str, default="data/token_counts.bert-base-uncased.pickle", help="The dump file."
       )
       parser.add_argument("--vocab_size", default=3_0522, type=int)
       UpperCAmelCase__  							=parser.parse_args()

       logger.info(f"""Loading data from {args.data_file}""")
       with open(args.data_file, "rb") as fp:
              UpperCAmelCase__  							=pickle.load(fp)

       logger.info("Counting occurrences for MLM.")
       UpperCAmelCase__  							=Counter()
       for tk_ids in data:
              counter.update(tk_ids)
       UpperCAmelCase__  							=[0] * args.vocab_size
       for k, v in counter.items():
              UpperCAmelCase__  							=v

       logger.info(f"""Dump to {args.token_counts_dump}""")
       with open(args.token_counts_dump, "wb") as handle:
              pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)
