


import unittest

from transformers import YosoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import (
			    YosoForMaskedLM,
			    YosoForMultipleChoice,
			    YosoForQuestionAnswering,
			    YosoForSequenceClassification,
			    YosoForTokenClassification,
			    YosoModel,
			)
			from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST


class 			__lowerCAmelCase				:


	'''simple docstring'''





	def __init__(	self:	Optional[Any],				lowerCamelCase_:	List[str],				lowerCamelCase_:	Union[str, Any]=13,				lowerCamelCase_:	Optional[int]=7,				lowerCamelCase_:	Any=True,				lowerCamelCase_:	Optional[int]=True,				lowerCamelCase_:	Tuple=True,				lowerCamelCase_:	List[str]=True,				lowerCamelCase_:	List[Any]=99,				lowerCamelCase_:	Union[str, Any]=32,				lowerCamelCase_:	str=5,				lowerCamelCase_:	Union[str, Any]=4,				lowerCamelCase_:	Tuple=37,				lowerCamelCase_:	List[str]="gelu",				lowerCamelCase_:	List[Any]=0.1,				lowerCamelCase_:	List[Any]=0.1,				lowerCamelCase_:	Optional[int]=512,				lowerCamelCase_:	Union[str, Any]=16,				lowerCamelCase_:	Tuple=2,				lowerCamelCase_:	List[Any]=0.0_2,				lowerCamelCase_:	Optional[Any]=3,				lowerCamelCase_:	Optional[Any]=4,				lowerCamelCase_:	Union[str, Any]=None,				):
			lowercase__				:			Optional[int]            =		parent
			lowercase__				:			Dict            =		batch_size
			lowercase__				:			List[Any]            =		seq_length
			lowercase__				:			int            =		is_training
			lowercase__				:			List[str]            =		use_input_mask
			lowercase__				:			Dict            =		use_token_type_ids
			lowercase__				:			Optional[Any]            =		use_labels
			lowercase__				:			Optional[int]            =		vocab_size
			lowercase__				:			Optional[Any]            =		hidden_size
			lowercase__				:			Optional[Any]            =		num_hidden_layers
			lowercase__				:			List[str]            =		num_attention_heads
			lowercase__				:			List[str]            =		intermediate_size
			lowercase__				:			Union[str, Any]            =		hidden_act
			lowercase__				:			int            =		hidden_dropout_prob
			lowercase__				:			Tuple            =		attention_probs_dropout_prob
			lowercase__				:			Optional[Any]            =		max_position_embeddings
			lowercase__				:			Union[str, Any]            =		type_vocab_size
			lowercase__				:			Tuple            =		type_sequence_label_size
			lowercase__				:			Optional[int]            =		initializer_range
			lowercase__				:			Any            =		num_labels
			lowercase__				:			Dict            =		num_choices
			lowercase__				:			Optional[Any]            =		scope

	def       snake_case__(	self:	Dict ):
			lowercase__				:			Any            =		ids_tensor([self.batch_size, self.seq_length],				self.vocab_size )

			lowercase__				:			List[str]            =		None
			if self.use_input_mask:
					lowercase__				:			Optional[int]            =		random_attention_mask([self.batch_size, self.seq_length] )

			lowercase__				:			List[str]            =		None
			if self.use_token_type_ids:
					lowercase__				:			List[Any]            =		ids_tensor([self.batch_size, self.seq_length],				self.type_vocab_size )

			lowercase__				:			int            =		None
			lowercase__				:			Tuple            =		None
			lowercase__				:			Optional[int]            =		None
			if self.use_labels:
					lowercase__				:			Optional[int]            =		ids_tensor([self.batch_size],				self.type_sequence_label_size )
					lowercase__				:			Optional[Any]            =		ids_tensor([self.batch_size, self.seq_length],				self.num_labels )
					lowercase__				:			Optional[int]            =		ids_tensor([self.batch_size],				self.num_choices )

			lowercase__				:			List[Any]            =		self.get_config()

			return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

	def       snake_case__(	self:	List[str] ):
			return YosoConfig(
			    vocab_size=self.vocab_size,				hidden_size=self.hidden_size,				num_hidden_layers=self.num_hidden_layers,				num_attention_heads=self.num_attention_heads,				intermediate_size=self.intermediate_size,				hidden_act=self.hidden_act,				hidden_dropout_prob=self.hidden_dropout_prob,				attention_probs_dropout_prob=self.attention_probs_dropout_prob,				max_position_embeddings=self.max_position_embeddings,				type_vocab_size=self.type_vocab_size,				is_decoder=lowerCamelCase_,				initializer_range=self.initializer_range,				)

	def       snake_case__(	self:	List[Any] ):
			lowercase__				:			Union[str, Any]            =		self.get_config()
			lowercase__				:			str            =		300
			return config

	def       snake_case__(	self:	List[Any] ):
			(
			    lowercase__							
			)				:			int            =		self.prepare_config_and_inputs()

			lowercase__				:			List[Any]            =		True
			lowercase__				:			Tuple            =		floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
			lowercase__				:			List[str]            =		ids_tensor([self.batch_size, self.seq_length],				vocab_size=2 )

			return (
			    config,
			    input_ids,
			    token_type_ids,
			    input_mask,
			    sequence_labels,
			    token_labels,
			    choice_labels,
			    encoder_hidden_states,
			    encoder_attention_mask,
			)

	def       snake_case__(	self:	Dict,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	Any,				lowerCamelCase_:	Dict,				lowerCamelCase_:	Dict,				lowerCamelCase_:	Optional[int],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	Dict ):
			lowercase__				:			int            =		YosoModel(config=lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			Any            =		model(lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_ )
			lowercase__				:			Tuple            =		model(lowerCamelCase_,				token_type_ids=lowerCamelCase_ )
			lowercase__				:			Union[str, Any]            =		model(lowerCamelCase_ )
			self.parent.assertEqual(result.last_hidden_state.shape,				(self.batch_size, self.seq_length, self.hidden_size) )

	def       snake_case__(	self:	List[Any],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	List[str],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	List[str],				lowerCamelCase_:	Dict,				lowerCamelCase_:	str,				lowerCamelCase_:	Tuple,				lowerCamelCase_:	List[str],				lowerCamelCase_:	str,				):
			lowercase__				:			List[Any]            =		True
			lowercase__				:			List[Any]            =		YosoModel(lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			Optional[int]            =		model(
			    lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				encoder_hidden_states=lowerCamelCase_,				encoder_attention_mask=lowerCamelCase_,				)
			lowercase__				:			Any            =		model(
			    lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				encoder_hidden_states=lowerCamelCase_,				)
			lowercase__				:			List[str]            =		model(lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_ )
			self.parent.assertEqual(result.last_hidden_state.shape,				(self.batch_size, self.seq_length, self.hidden_size) )

	def       snake_case__(	self:	List[Any],				lowerCamelCase_:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	str,				lowerCamelCase_:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	str,				lowerCamelCase_:	Optional[int] ):
			lowercase__				:			str            =		YosoForMaskedLM(config=lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			str            =		model(lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				labels=lowerCamelCase_ )
			self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.seq_length, self.vocab_size) )

	def       snake_case__(	self:	Dict,				lowerCamelCase_:	Union[str, Any],				lowerCamelCase_:	Any,				lowerCamelCase_:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	str,				lowerCamelCase_:	Optional[int],				lowerCamelCase_:	Dict ):
			lowercase__				:			Any            =		YosoForQuestionAnswering(config=lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			Optional[Any]            =		model(
			    lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				start_positions=lowerCamelCase_,				end_positions=lowerCamelCase_,				)
			self.parent.assertEqual(result.start_logits.shape,				(self.batch_size, self.seq_length) )
			self.parent.assertEqual(result.end_logits.shape,				(self.batch_size, self.seq_length) )

	def       snake_case__(	self:	Dict,				lowerCamelCase_:	List[str],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	Tuple,				lowerCamelCase_:	List[Any],				lowerCamelCase_:	Any ):
			lowercase__				:			int            =		self.num_labels
			lowercase__				:			List[str]            =		YosoForSequenceClassification(lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			int            =		model(lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				labels=lowerCamelCase_ )
			self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.num_labels) )

	def       snake_case__(	self:	Optional[Any],				lowerCamelCase_:	Any,				lowerCamelCase_:	Dict,				lowerCamelCase_:	Any,				lowerCamelCase_:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	List[Any],				lowerCamelCase_:	int ):
			lowercase__				:			Optional[Any]            =		self.num_labels
			lowercase__				:			Optional[Any]            =		YosoForTokenClassification(config=lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			str            =		model(lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				labels=lowerCamelCase_ )
			self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.seq_length, self.num_labels) )

	def       snake_case__(	self:	int,				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	Optional[int],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	List[str],				lowerCamelCase_:	Optional[int],				lowerCamelCase_:	Optional[Any],				lowerCamelCase_:	Any ):
			lowercase__				:			List[Any]            =		self.num_choices
			lowercase__				:			Union[str, Any]            =		YosoForMultipleChoice(config=lowerCamelCase_ )
			model.to(lowerCamelCase_ )
			model.eval()
			lowercase__				:			Tuple            =		input_ids.unsqueeze(1 ).expand(-1,				self.num_choices,				-1 ).contiguous()
			lowercase__				:			List[str]            =		token_type_ids.unsqueeze(1 ).expand(-1,				self.num_choices,				-1 ).contiguous()
			lowercase__				:			Any            =		input_mask.unsqueeze(1 ).expand(-1,				self.num_choices,				-1 ).contiguous()
			lowercase__				:			int            =		model(
			    lowerCamelCase_,				attention_mask=lowerCamelCase_,				token_type_ids=lowerCamelCase_,				labels=lowerCamelCase_,				)
			self.parent.assertEqual(result.logits.shape,				(self.batch_size, self.num_choices) )





	def       snake_case__(	self:	Optional[int] ):
			lowercase__				:			Union[str, Any]            =		self.prepare_config_and_inputs()
			(
			    lowercase__							
			)				:			Optional[int]            =		config_and_inputs
			lowercase__				:			Tuple            =		{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
			return config, inputs_dict


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


	'''simple docstring'''





	_A       	=							(
	    (
	        YosoModel,
	        YosoForMaskedLM,
	        YosoForMultipleChoice,
	        YosoForQuestionAnswering,
	        YosoForSequenceClassification,
	        YosoForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_A       	=							False
	_A       	=							False
	_A       	=							False

	_A       	=							()
	_A       	=							(
	    {
	        "feature-extraction": YosoModel,
	        "fill-mask": YosoForMaskedLM,
	        "question-answering": YosoForQuestionAnswering,
	        "text-classification": YosoForSequenceClassification,
	        "token-classification": YosoForTokenClassification,
	        "zero-shot": YosoForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)

	def       snake_case__(	self:	Dict ):
			lowercase__				:			List[str]            =		YosoModelTester(self )
			lowercase__				:			Tuple            =		ConfigTester(self,				config_class=lowerCamelCase_,				hidden_size=37 )

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

	def       snake_case__(	self:	Union[str, Any] ):
			lowercase__				:			Tuple            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*lowerCamelCase_ )

	def       snake_case__(	self:	Dict ):
			lowercase__				:			Dict            =		self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
					lowercase__				:			Any            =		type
					self.model_tester.create_and_check_model(*lowerCamelCase_ )

	def       snake_case__(	self:	Union[str, Any] ):
			lowercase__				:			Union[str, Any]            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase_ )

	def       snake_case__(	self:	Optional[int] ):
			lowercase__				:			Union[str, Any]            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase_ )

	def       snake_case__(	self:	Tuple ):
			lowercase__				:			Union[str, Any]            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_question_answering(*lowerCamelCase_ )

	def       snake_case__(	self:	Optional[Any] ):
			lowercase__				:			Tuple            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase_ )

	def       snake_case__(	self:	List[Any] ):
			lowercase__				:			Optional[int]            =		self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_token_classification(*lowerCamelCase_ )

	@slow
	def       snake_case__(	self:	Optional[Any] ):
			for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowercase__				:			Dict            =		YosoModel.from_pretrained(lowerCamelCase_ )
					self.assertIsNotNone(lowerCamelCase_ )





	def       snake_case__(	self:	Dict ):
			return


@require_torch
class 			__lowerCAmelCase				(					unittest.TestCase			):


	'''simple docstring'''





	@slow
	def       snake_case__(	self:	List[str] ):
			lowercase__				:			Optional[int]            =		YosoModel.from_pretrained('uw-madison/yoso-4096' )
			lowercase__				:			Optional[Any]            =		torch.tensor([[0, 1, 2, 3, 4, 5]] )

			with torch.no_grad():
					lowercase__				:			Tuple            =		model(lowerCamelCase_ )[0]

			lowercase__				:			Union[str, Any]            =		torch.Size((1, 6, 768) )
			self.assertEqual(output.shape,				lowerCamelCase_ )

			lowercase__				:			Optional[int]            =		torch.tensor(
			    [[[-0.0_6_1_1, 0.1_2_4_2, 0.0_8_4_0], [0.0_2_8_0, -0.0_0_4_8, 0.1_1_2_5], [0.0_1_0_6, 0.0_2_2_6, 0.0_7_5_1]]] )

			self.assertTrue(torch.allclose(output[:, :3, :3],				lowerCamelCase_,				atol=1E-4 ) )

	@slow
	def       snake_case__(	self:	Optional[Any] ):
			lowercase__				:			int            =		YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096' )
			lowercase__				:			int            =		torch.tensor([[0, 1, 2, 3, 4, 5]] )

			with torch.no_grad():
					lowercase__				:			List[Any]            =		model(lowerCamelCase_ )[0]

			lowercase__				:			Union[str, Any]            =		50265

			lowercase__				:			Union[str, Any]            =		torch.Size((1, 6, vocab_size) )
			self.assertEqual(output.shape,				lowerCamelCase_ )

			lowercase__				:			Dict            =		torch.tensor(
			    [[[-2.1_3_1_3, -3.7_2_8_5, -2.2_4_0_7], [-2.7_0_4_7, -3.3_3_1_4, -2.6_4_0_8], [0.0_6_2_9, -2.5_1_6_6, -0.3_3_5_6]]] )

			self.assertTrue(torch.allclose(output[:, :3, :3],				lowerCamelCase_,				atol=1E-4 ) )





	@slow
	def       snake_case__(	self:	Dict ):
			lowercase__				:			Tuple            =		YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096' )
			lowercase__				:			int            =		torch.arange(4096 ).unsqueeze(0 )

			with torch.no_grad():
					lowercase__				:			List[str]            =		model(lowerCamelCase_ )[0]

			lowercase__				:			List[str]            =		50265

			lowercase__				:			Optional[int]            =		torch.Size((1, 4096, vocab_size) )
			self.assertEqual(output.shape,				lowerCamelCase_ )

			lowercase__				:			List[Any]            =		torch.tensor(
			    [[[-2.3_9_1_4, -4.3_7_4_2, -5.0_9_5_6], [-4.0_9_8_8, -4.2_3_8_4, -7.0_4_0_6], [-3.1_4_2_7, -3.7_1_9_2, -6.6_8_0_0]]] )

			self.assertTrue(torch.allclose(output[:, :3, :3],				lowerCamelCase_,				atol=1E-4 ) )









from typing import List

import PIL
import torch
from PIL import Image

from ...configuration_utils import ConfigMixin
from ...models.modeling_utils import ModelMixin
from ...utils import PIL_INTERPOLATION


class 			__lowerCAmelCase				(					_UpperCamelCase      ,						_UpperCamelCase			):


	'''simple docstring'''





	def __init__(	self:	Union[str, Any] ):
			super().__init__()

			self.register_buffer('watermark_image',				torch.zeros((62, 62, 4) ) )
			lowercase__				:			int            =		None





	def       snake_case__(	self:	Optional[int],				lowerCamelCase_:	List[PIL.Image.Image],				lowerCamelCase_:	Dict=None ):
			# copied from https://github.com/deep-floyd/IF/blob/b77482e36ca2031cb94dbca1001fc1e6400bf4ab/deepfloyd_if/modules/base.py#L287

			lowercase__				:			str            =		images[0].height
			lowercase__				:			Optional[Any]            =		images[0].width

			lowercase__				:			List[Any]            =		sample_size or h

			lowercase__				:			Any            =		min(h / sample_size,				w / sample_size )
			lowercase__      ,							lowercase__				:			Optional[int]            =		(int(h / coef ), int(w / coef )) if coef < 1 else (h, w)

			lowercase__      ,							lowercase__				:			Any            =		1024**2, img_w * img_h
			lowercase__				:			Union[str, Any]            =		(Sa / Sa) ** 0.5
			lowercase__      ,							lowercase__      ,							lowercase__				:			Optional[int]            =		int(K * 62 ), img_w - int(14 * K ), img_h - int(14 * K )

			if self.watermark_image_as_pil is None:
					lowercase__				:			Tuple            =		self.watermark_image.to(torch.uinta ).cpu().numpy()
					lowercase__				:			str            =		Image.fromarray(lowerCamelCase_,				mode='RGBA' )
					lowercase__				:			Tuple            =		watermark_image

			lowercase__				:			Optional[int]            =		self.watermark_image_as_pil.resize(
			    (wm_size, wm_size),				PIL_INTERPOLATION['bicubic'],				reducing_gap=lowerCamelCase_ )

			for pil_img in images:
					pil_img.paste(lowerCamelCase_,				box=(wm_x - wm_size, wm_y - wm_size, wm_x, wm_y),				mask=wm_img.split()[-1] )

			return images


