






'''simple docstring'''







import inspect
import unittest

import requests

from transformers import MgpstrConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

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


if is_torch_available():
  import torch
  from torch import nn

  from transformers import MgpstrForSceneTextRecognition


if is_vision_available():
  from PIL import Image

  from transformers import MgpstrProcessor




class 					__UpperCAmelCase :

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase=False  ,		_UpperCAmelCase=13  ,		_UpperCAmelCase=(32, 128)  ,		_UpperCAmelCase=4  ,		_UpperCAmelCase=3  ,		_UpperCAmelCase=27  ,		_UpperCAmelCase=38  ,		_UpperCAmelCase=99  ,		_UpperCAmelCase=99  ,		_UpperCAmelCase=32  ,		_UpperCAmelCase=5  ,		_UpperCAmelCase=4  ,		_UpperCAmelCase=4.0  ,		_UpperCAmelCase=257  ,		_UpperCAmelCase=None  ,		):
             UpperCAmelCase__	:   Union[str, Any]												=	parent
             UpperCAmelCase__	:   Union[str, Any]												=	is_training
             UpperCAmelCase__	:   str												=	batch_size
             UpperCAmelCase__	:   int												=	image_size
             UpperCAmelCase__	:   Optional[int]												=	patch_size
             UpperCAmelCase__	:   List[str]												=	num_channels
             UpperCAmelCase__	:   int												=	max_token_length
             UpperCAmelCase__	:   Dict												=	num_character_labels
             UpperCAmelCase__	:   Optional[Any]												=	num_bpe_labels
             UpperCAmelCase__	:   List[Any]												=	num_wordpiece_labels
             UpperCAmelCase__	:   List[str]												=	hidden_size
             UpperCAmelCase__	:   Optional[Any]												=	num_hidden_layers
             UpperCAmelCase__	:   str												=	num_attention_heads
             UpperCAmelCase__	:   List[str]												=	mlp_ratio
             UpperCAmelCase__	:   str												=	patch_embeds_hidden_size
             UpperCAmelCase__	:   List[Any]												=	output_hidden_states




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   Dict												=	floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]						)
             UpperCAmelCase__	:   Any												=	self.get_config()
             return config, pixel_values




      def 							lowerCamelCase     (					self						):
             return MgpstrConfig(
                 image_size=self.image_size  ,		patch_size=self.patch_size  ,		num_channels=self.num_channels  ,		max_token_length=self.max_token_length  ,		num_character_labels=self.num_character_labels  ,		num_bpe_labels=self.num_bpe_labels  ,		num_wordpiece_labels=self.num_wordpiece_labels  ,		hidden_size=self.hidden_size  ,		num_hidden_layers=self.num_hidden_layers  ,		num_attention_heads=self.num_attention_heads  ,		mlp_ratio=self.mlp_ratio  ,		output_hidden_states=self.output_hidden_states  ,		)




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase						):
             UpperCAmelCase__	:   int												=	MgpstrForSceneTextRecognition(_UpperCAmelCase						)
             model.to(_UpperCAmelCase						)
             model.eval()
             with torch.no_grad():
                    UpperCAmelCase__	:   List[Any]												=	model(_UpperCAmelCase						)
             self.parent.assertEqual(
                 generated_ids[0][0].shape  ,		(self.batch_size, self.max_token_length, self.num_character_labels)						)




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   List[Any]												=	self.prepare_config_and_inputs()
             UpperCAmelCase__							,							UpperCAmelCase__	:   Tuple												=	config_and_inputs
             UpperCAmelCase__	:   List[str]												=	{'''pixel_values''': pixel_values}
             return config, inputs_dict




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

      '''simple docstring'''





      SCREAMING_SNAKE_CASE				:    Optional[int]				       =					(MgpstrForSceneTextRecognition,) if is_torch_available() else ()
      SCREAMING_SNAKE_CASE				:    Optional[int]				       =					{"feature-extraction": MgpstrForSceneTextRecognition} if is_torch_available() else {}
      SCREAMING_SNAKE_CASE				:    Optional[int]				       =					False

      SCREAMING_SNAKE_CASE				:    List[Any]				       =					False
      SCREAMING_SNAKE_CASE				:    List[str]				       =					False
      SCREAMING_SNAKE_CASE				:    List[str]				       =					False
      SCREAMING_SNAKE_CASE				:    Union[str, Any]				       =					False




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   Dict												=	MgpstrModelTester(self						)
             UpperCAmelCase__	:   List[Any]												=	ConfigTester(self  ,		config_class=_UpperCAmelCase  ,		has_text_modality=_UpperCAmelCase						)




      def 							lowerCamelCase     (					self						):
             self.config_tester.run_common_tests()




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   Union[str, Any]												=	self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCAmelCase						)




      @unittest.skip(reason='''MgpstrModel does not use inputs_embeds'''						)
      def 							lowerCamelCase     (					self						):
             pass




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__							,							UpperCAmelCase__	:   Dict												=	self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    UpperCAmelCase__	:   Union[str, Any]												=	model_class(_UpperCAmelCase						)
                    self.assertIsInstance(model.get_input_embeddings()  ,		(nn.Module)						)
                    UpperCAmelCase__	:   Optional[int]												=	model.get_output_embeddings()
                    self.assertTrue(x is None or isinstance(_UpperCAmelCase  ,		nn.Linear						)						)




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__							,							UpperCAmelCase__	:   Optional[int]												=	self.model_tester.prepare_config_and_inputs_for_common()

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

                    UpperCAmelCase__	:   Any												=	['''pixel_values''']
                    self.assertListEqual(arg_names[:1]  ,		_UpperCAmelCase						)




      @unittest.skip(reason='''MgpstrModel does not support feedforward chunking'''						)
      def 							lowerCamelCase     (					self						):
             pass




      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__							,							UpperCAmelCase__	:   int												=	self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    if not model_class.supports_gradient_checkpointing:
                           continue

                    UpperCAmelCase__	:   List[str]												=	True
                    UpperCAmelCase__	:   Dict												=	model_class(_UpperCAmelCase						)
                    self.assertTrue(model.is_gradient_checkpointing						)




      def 							lowerCamelCase     (					self						):
             def check_hidden_states_output(_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase						):
                    UpperCAmelCase__	:   Dict												=	model_class(_UpperCAmelCase						)
                    model.to(_UpperCAmelCase						)
                    model.eval()

                    with torch.no_grad():
                           UpperCAmelCase__	:   Any												=	model(**self._prepare_for_class(_UpperCAmelCase  ,		_UpperCAmelCase						)						)

                    UpperCAmelCase__	:   Any												=	outputs.hidden_states

                    UpperCAmelCase__	:   List[str]												=	getattr(
                        self.model_tester  ,		'''expected_num_hidden_layers'''  ,		self.model_tester.num_hidden_layers + 1						)
                    self.assertEqual(len(_UpperCAmelCase						)  ,		_UpperCAmelCase						)

                    self.assertListEqual(
                        list(hidden_states[0].shape[-2:]						)  ,		[self.model_tester.patch_embeds_hidden_size, self.model_tester.hidden_size]  ,		)

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

             for model_class in self.all_model_classes:
                    UpperCAmelCase__	:   List[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						):
             UpperCAmelCase__							,							UpperCAmelCase__	:   Tuple												=	self.model_tester.prepare_config_and_inputs_for_common()

             UpperCAmelCase__	:   int												=	_config_zero_init(_UpperCAmelCase						)
             for model_class in self.all_model_classes:
                    UpperCAmelCase__	:   List[str]												=	model_class(config=_UpperCAmelCase						)
                    for name, param in model.named_parameters():
                           if isinstance(_UpperCAmelCase  ,		(nn.Linear, nn.Convad, nn.LayerNorm)						):
                                  if param.requires_grad:
                                         self.assertIn(
                                             ((param.data.mean() * 1E9).round() / 1E9).item()  ,		[0.0, 1.0]  ,		msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""  ,		)




      @unittest.skip(reason='''Retain_grad is tested in individual model tests'''						)
      def 							lowerCamelCase     (					self						):
             pass






def  lowerCAmelCase__		(					)   -> List[Any]:
       UpperCAmelCase__	:   str												=	'''https://i.postimg.cc/ZKwLg2Gw/367-14.png'''
       UpperCAmelCase__	:   Tuple												=	Image.open(requests.get(a_			,						stream=a_			).raw			).convert('''RGB'''			)
       return im




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

      '''simple docstring'''





      @slow
      def 							lowerCamelCase     (					self						):
             UpperCAmelCase__	:   Optional[Any]												=	'''alibaba-damo/mgp-str-base'''
             UpperCAmelCase__	:   Tuple												=	MgpstrForSceneTextRecognition.from_pretrained(_UpperCAmelCase						).to(_UpperCAmelCase						)
             UpperCAmelCase__	:   Any												=	MgpstrProcessor.from_pretrained(_UpperCAmelCase						)

             UpperCAmelCase__	:   List[Any]												=	prepare_img()
             UpperCAmelCase__	:   Optional[Any]												=	processor(images=_UpperCAmelCase  ,		return_tensors='''pt'''						).pixel_values.to(_UpperCAmelCase						)

             # forward pass
             with torch.no_grad():
                    UpperCAmelCase__	:   Dict												=	model(_UpperCAmelCase						)

             # verify the logits
             self.assertEqual(outputs.logits[0].shape  ,		torch.Size((1, 27, 38)						)						)

             UpperCAmelCase__	:   List[Any]												=	processor.batch_decode(outputs.logits						)
             UpperCAmelCase__	:   Optional[int]												=	'''ticket'''

             self.assertEqual(out_strs['''generated_text'''][0]  ,		_UpperCAmelCase						)

             UpperCAmelCase__	:   List[Any]												=	torch.tensor(
                 [[[-39.53_97, -44.40_24, -36.18_44], [-61.47_09, -63.86_39, -58.34_54], [-74.02_25, -68.54_94, -71.21_64]]]  ,		device=_UpperCAmelCase  ,		)

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







'''simple docstring'''







import math

import torch
from torch import nn
from torch.nn import init

from .binarizer import MagnitudeBinarizer, ThresholdBinarizer, TopKBinarizer




class 					__UpperCAmelCase (		nn.Linear	):

      '''simple docstring'''





      def __init__(					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase  ,		_UpperCAmelCase = True  ,		_UpperCAmelCase = "constant"  ,		_UpperCAmelCase = 0.0  ,		_UpperCAmelCase = "topK"  ,		):
             super(_UpperCAmelCase  ,		self						).__init__(in_features=_UpperCAmelCase  ,		out_features=_UpperCAmelCase  ,		bias=_UpperCAmelCase						)
             assert pruning_method in ["topK", "threshold", "sigmoied_threshold", "magnitude", "l0"]
             UpperCAmelCase__	:   Union[str, Any]												=	pruning_method

             if self.pruning_method in ["topK", "threshold", "sigmoied_threshold", "l0"]:
                    UpperCAmelCase__	:   Union[str, Any]												=	mask_scale
                    UpperCAmelCase__	:   Optional[int]												=	mask_init
                    UpperCAmelCase__	:   List[str]												=	nn.Parameter(torch.empty(self.weight.size()						)						)
                    self.init_mask()




      def 							lowerCamelCase     (					self						):
             if self.mask_init == "constant":
                    init.constant_(self.mask_scores  ,		val=self.mask_scale						)
             elif self.mask_init == "uniform":
                    init.uniform_(self.mask_scores  ,		a=-self.mask_scale  ,		b=self.mask_scale						)
             elif self.mask_init == "kaiming":
                    init.kaiming_uniform_(self.mask_scores  ,		a=math.sqrt(5						)						)




      def 							lowerCamelCase     (					self  ,		_UpperCAmelCase  ,		_UpperCAmelCase						):
             # Get the mask
             if self.pruning_method == "topK":
                    UpperCAmelCase__	:   Union[str, Any]												=	TopKBinarizer.apply(self.mask_scores  ,		_UpperCAmelCase						)
             elif self.pruning_method in ["threshold", "sigmoied_threshold"]:
                    UpperCAmelCase__	:   List[Any]												=	'''sigmoied''' in self.pruning_method
                    UpperCAmelCase__	:   Tuple												=	ThresholdBinarizer.apply(self.mask_scores  ,		_UpperCAmelCase  ,		_UpperCAmelCase						)
             elif self.pruning_method == "magnitude":
                    UpperCAmelCase__	:   Optional[int]												=	MagnitudeBinarizer.apply(self.weight  ,		_UpperCAmelCase						)
             elif self.pruning_method == "l0":
                    UpperCAmelCase__							,							UpperCAmelCase__							,							UpperCAmelCase__	:   int												=	-0.1, 1.1, 2 / 3
                    if self.training:
                           UpperCAmelCase__	:   Optional[int]												=	torch.zeros_like(self.mask_scores						).uniform_().clamp(0.00_01  ,		0.99_99						)
                           UpperCAmelCase__	:   Tuple												=	torch.sigmoid((u.log() - (1 - u).log() + self.mask_scores) / b						)
                    else:
                           UpperCAmelCase__	:   Any												=	torch.sigmoid(self.mask_scores						)
                    UpperCAmelCase__	:   List[Any]												=	s * (r - l) + l
                    UpperCAmelCase__	:   Optional[int]												=	s_bar.clamp(min=0.0  ,		max=1.0						)
             # Mask weights with computed mask
             UpperCAmelCase__	:   Union[str, Any]												=	mask * self.weight
             # Compute output (linear layer) with masked weights
             return nn.functional.linear(_UpperCAmelCase  ,		_UpperCAmelCase  ,		self.bias						)