



"""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, 1_28)							,UpperCAmelCase_=4							,UpperCAmelCase_=3							,UpperCAmelCase_=27							,UpperCAmelCase_=38							,UpperCAmelCase_=99							,UpperCAmelCase_=99							,UpperCAmelCase_=32							,UpperCAmelCase_=5							,UpperCAmelCase_=4							,UpperCAmelCase_=4.0							,UpperCAmelCase_=2_57							,UpperCAmelCase_=None							,):
			_lowercase							:      Optional[int]			=					parent
			_lowercase							:      int			=					is_training
			_lowercase							:      Dict			=					batch_size
			_lowercase							:      List[str]			=					image_size
			_lowercase							:      Optional[int]			=					patch_size
			_lowercase							:      List[str]			=					num_channels
			_lowercase							:      int			=					max_token_length
			_lowercase							:      Dict			=					num_character_labels
			_lowercase							:      Tuple			=					num_bpe_labels
			_lowercase							:      str			=					num_wordpiece_labels
			_lowercase							:      str			=					hidden_size
			_lowercase							:      Tuple			=					num_hidden_layers
			_lowercase							:      int			=					num_attention_heads
			_lowercase							:      Optional[Any]			=					mlp_ratio
			_lowercase							:      Union[str, Any]			=					patch_embeds_hidden_size
			_lowercase							:      int			=					output_hidden_states






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      Tuple			=					floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]					)
			_lowercase							:      Optional[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_					):
			_lowercase							:      Optional[int]			=					MgpstrForSceneTextRecognition(UpperCAmelCase_					)
			model.to(UpperCAmelCase_					)
			model.eval()
			with torch.no_grad():
					_lowercase							:      Union[str, 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					):
			_lowercase							:      str			=					self.prepare_config_and_inputs()
			_lowercase			,   _lowercase							:      Tuple			=					config_and_inputs
			_lowercase							:      List[Any]			=					{"""pixel_values""": pixel_values}
			return config, inputs_dict







@require_torch
class    UpperCamelCase     (			snake_case ,					snake_case ,					unittest.TestCase    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    Any		    =   (MgpstrForSceneTextRecognition,) if is_torch_available() else ()
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   {"feature-extraction": MgpstrForSceneTextRecognition} if is_torch_available() else {}
	SCREAMING_SNAKE_CASE_				:    Optional[Any]		    =   False

	SCREAMING_SNAKE_CASE_				:    str		    =   False
	SCREAMING_SNAKE_CASE_				:    str		    =   False
	SCREAMING_SNAKE_CASE_				:    Any		    =   False
	SCREAMING_SNAKE_CASE_				:    int		    =   False






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      str			=					MgpstrModelTester(self					)
			_lowercase							:      List[str]			=					ConfigTester(self							,config_class=UpperCAmelCase_							,has_text_modality=UpperCAmelCase_					)






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






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      int			=					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					):
			_lowercase			,   _lowercase							:      str			=					self.model_tester.prepare_config_and_inputs_for_common()

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






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Optional[Any]			=					self.model_tester.prepare_config_and_inputs_for_common()

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

					_lowercase							:      Dict			=					["""pixel_values"""]
					self.assertListEqual(arg_names[:1]							,UpperCAmelCase_					)






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






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Dict			=					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

					_lowercase							:      List[str]			=					True
					_lowercase							:      List[str]			=					model_class(UpperCAmelCase_					)
					self.assertTrue(model.is_gradient_checkpointing					)






	def 					lowerCamelCase__ (			self					):
			def check_hidden_states_output(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_					):
					_lowercase							:      Optional[int]			=					model_class(UpperCAmelCase_					)
					model.to(UpperCAmelCase_					)
					model.eval()

					with torch.no_grad():
							_lowercase							:      Union[str, Any]			=					model(**self._prepare_for_class(UpperCAmelCase_							,UpperCAmelCase_					)					)

					_lowercase							:      str			=					outputs.hidden_states

					_lowercase							:      Optional[int]			=					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]							,)

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

			for model_class in self.all_model_classes:
					_lowercase							:      Tuple			=					True
					check_hidden_states_output(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_					)

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

					check_hidden_states_output(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Optional[int]			=					self.model_tester.prepare_config_and_inputs_for_common()

			_lowercase							:      List[Any]			=					_config_zero_init(UpperCAmelCase_					)
			for model_class in self.all_model_classes:
					_lowercase							:      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       __SCREAMING_SNAKE_CASE      (   ):
		_lowercase							:      Tuple			=					"""https://i.postimg.cc/ZKwLg2Gw/367-14.png"""
		_lowercase							:      Optional[Any]			=					Image.open(requests.get(__UpperCAmelCase					,							stream=__UpperCAmelCase	).raw	).convert("""RGB"""	)
		return im





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




	"""simple docstring"""






	@slow
	def 					lowerCamelCase__ (			self					):
			_lowercase							:      List[str]			=					"""alibaba-damo/mgp-str-base"""
			_lowercase							:      str			=					MgpstrForSceneTextRecognition.from_pretrained(UpperCAmelCase_					).to(UpperCAmelCase_					)
			_lowercase							:      Optional[int]			=					MgpstrProcessor.from_pretrained(UpperCAmelCase_					)

			_lowercase							:      Any			=					prepare_img()
			_lowercase							:      Dict			=					processor(images=UpperCAmelCase_							,return_tensors="""pt"""					).pixel_values.to(UpperCAmelCase_					)

			# forward pass
			with torch.no_grad():
					_lowercase							:      int			=					model(UpperCAmelCase_					)

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

			_lowercase							:      List[str]			=					processor.batch_decode(outputs.logits					)
			_lowercase							:      Tuple			=					"""ticket"""

			self.assertEqual(out_strs["""generated_text"""][0]							,UpperCAmelCase_					)

			_lowercase							:      Optional[int]			=					torch.tensor(
			    [[[-39.5397, -44.4024, -36.1844], [-61.4709, -63.8639, -58.3454], [-74.0225, -68.5494, -71.2164]]]							,device=UpperCAmelCase_							,)

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








"""simple docstring"""







import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


UpperCAmelCase:   Dict       			=							logging.get_logger(__name__)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      List[str]			=					nn.functional.normalize(__UpperCAmelCase	)
		_lowercase							:      int			=					nn.functional.normalize(__UpperCAmelCase	)
		return torch.mm(__UpperCAmelCase					,							normalized_text_embeds.t()	)





class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   CLIPConfig

	SCREAMING_SNAKE_CASE_				:    List[Any]		    =   ["CLIPEncoderLayer"]






	def __init__(			self							,UpperCAmelCase_					):
			super().__init__(UpperCAmelCase_					)

			_lowercase							:      Any			=					CLIPVisionModel(config.vision_config					)
			_lowercase							:      Any			=					nn.Linear(config.vision_config.hidden_size							,config.projection_dim							,bias=UpperCAmelCase_					)

			_lowercase							:      Dict			=					nn.Parameter(torch.ones(17							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      Tuple			=					nn.Parameter(torch.ones(3							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)

			_lowercase							:      Optional[int]			=					nn.Parameter(torch.ones(17					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					nn.Parameter(torch.ones(3					)							,requires_grad=UpperCAmelCase_					)






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[str]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      str			=					self.visual_projection(UpperCAmelCase_					)

			# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					).cpu().float().numpy()
			_lowercase							:      str			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					).cpu().float().numpy()

			_lowercase							:      List[str]			=					[]
			_lowercase							:      Optional[Any]			=					image_embeds.shape[0]
			for i in range(UpperCAmelCase_					):
					_lowercase							:      int			=					{"""special_scores""": {}, """special_care""": [], """concept_scores""": {}, """bad_concepts""": []}

					# increase this value to create a stronger `nfsw` filter
					# at the cost of increasing the possibility of filtering benign images
					_lowercase							:      str			=					0.0

					for concept_idx in range(len(special_cos_dist[0]					)					):
							_lowercase							:      Optional[Any]			=					special_cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.special_care_embeds_weights[concept_idx].item()
							_lowercase							:      Tuple			=					round(concept_cos - concept_threshold + adjustment							,3					)
							if result_img["special_scores"][concept_idx] > 0:
									result_img["special_care"].append({concept_idx, result_img["""special_scores"""][concept_idx]}					)
									_lowercase							:      Union[str, Any]			=					0.01

					for concept_idx in range(len(cos_dist[0]					)					):
							_lowercase							:      Union[str, Any]			=					cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.concept_embeds_weights[concept_idx].item()
							_lowercase							:      str			=					round(concept_cos - concept_threshold + adjustment							,3					)
							if result_img["concept_scores"][concept_idx] > 0:
									result_img["bad_concepts"].append(UpperCAmelCase_					)

					result.append(UpperCAmelCase_					)

			_lowercase							:      List[str]			=					[len(res["""bad_concepts"""]					) > 0 for res in result]

			for idx, has_nsfw_concept in enumerate(UpperCAmelCase_					):
					if has_nsfw_concept:
							if torch.is_tensor(UpperCAmelCase_					) or torch.is_tensor(images[0]					):
									_lowercase							:      List[Any]			=					torch.zeros_like(images[idx]					)  # black image
							else:
									_lowercase							:      Any			=					np.zeros(images[idx].shape					)  # black image

			if any(UpperCAmelCase_					):
					logger.warning(
					    """Potential NSFW content was detected in one or more images. A black image will be returned instead."""
					    """ Try again with a different prompt and/or seed."""					)

			return images, has_nsfw_concepts






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[Any]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      Optional[int]			=					self.visual_projection(UpperCAmelCase_					)

			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					)
			_lowercase							:      Tuple			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					)

			# increase this value to create a stronger `nsfw` filter
			# at the cost of increasing the possibility of filtering benign images
			_lowercase							:      Any			=					0.0

			_lowercase							:      Tuple			=					special_cos_dist - self.special_care_embeds_weights + adjustment
			# special_scores = special_scores.round(decimals=3)
			_lowercase							:      Dict			=					torch.any(special_scores > 0							,dim=1					)
			_lowercase							:      Any			=					special_care * 0.01
			_lowercase							:      Optional[int]			=					special_adjustment.unsqueeze(1					).expand(-1							,cos_dist.shape[1]					)

			_lowercase							:      Optional[Any]			=					(cos_dist - self.concept_embeds_weights) + special_adjustment
			# concept_scores = concept_scores.round(decimals=3)
			_lowercase							:      Optional[int]			=					torch.any(concept_scores > 0							,dim=1					)

			_lowercase							:      Tuple			=					0.0  # black image

			return images, has_nsfw_concepts



