
'''simple docstring'''
import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
						import torch
						from torch import nn

						from transformers import DinatBackbone, DinatForImageClassification, DinatModel
						from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor


class        SCREAMING_SNAKE_CASE__	:







				def __init__(    self		,							A_		,							A_=13		,							A_=64		,							A_=4		,							A_=3		,							A_=16		,							A_=[1, 2, 1]		,							A_=[2, 4, 8]		,							A_=3		,							A_=[[3], [1, 2], [1]]		,							A_=2.0		,							A_=True		,							A_=0.0		,							A_=0.0		,							A_=0.1		,							A_="gelu"		,							A_=True		,							A_=0.02		,							A_=1e-5		,							A_=True		,							A_=None		,							A_=True		,							A_=10		,							A_=["stage1", "stage2"]		,							A_=[1, 2]		,							)->						List[str]:







											'''simple docstring'''






											UpperCamelCase					       =			parent
											UpperCamelCase					       =			batch_size
											UpperCamelCase					       =			image_size
											UpperCamelCase					       =			patch_size
											UpperCamelCase					       =			num_channels
											UpperCamelCase					       =			embed_dim
											UpperCamelCase					       =			depths
											UpperCamelCase					       =			num_heads
											UpperCamelCase					       =			kernel_size
											UpperCamelCase					       =			dilations
											UpperCamelCase					       =			mlp_ratio
											UpperCamelCase					       =			qkv_bias
											UpperCamelCase					       =			hidden_dropout_prob
											UpperCamelCase					       =			attention_probs_dropout_prob
											UpperCamelCase					       =			drop_path_rate
											UpperCamelCase					       =			hidden_act
											UpperCamelCase					       =			patch_norm
											UpperCamelCase					       =			layer_norm_eps
											UpperCamelCase					       =			initializer_range
											UpperCamelCase					       =			is_training
											UpperCamelCase					       =			scope
											UpperCamelCase					       =			use_labels
											UpperCamelCase					       =			num_labels
											UpperCamelCase					       =			out_features
											UpperCamelCase					       =			out_indices







				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											UpperCamelCase					       =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

											UpperCamelCase					       =			self.get_config()

											return config, pixel_values, labels







				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											return DinatConfig(
											    num_labels=self.num_labels		,							image_size=self.image_size		,							patch_size=self.patch_size		,							num_channels=self.num_channels		,							embed_dim=self.embed_dim		,							depths=self.depths		,							num_heads=self.num_heads		,							kernel_size=self.kernel_size		,							dilations=self.dilations		,							mlp_ratio=self.mlp_ratio		,							qkv_bias=self.qkv_bias		,							hidden_dropout_prob=self.hidden_dropout_prob		,							attention_probs_dropout_prob=self.attention_probs_dropout_prob		,							drop_path_rate=self.drop_path_rate		,							hidden_act=self.hidden_act		,							patch_norm=self.patch_norm		,							layer_norm_eps=self.layer_norm_eps		,							initializer_range=self.initializer_range		,							out_features=self.out_features		,							out_indices=self.out_indices		,							)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							A_			)->						Optional[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			DinatModel(config=A_			)
											model.to(A_			)
											model.eval()
											UpperCamelCase					       =			model(A_			)

											UpperCamelCase					       =			UpperCamelCase					       =			(config.image_size // config.patch_size) // (2 ** (len(config.depths			) - 1))
											UpperCamelCase					       =			int(config.embed_dim * 2 ** (len(config.depths			) - 1)			)

											self.parent.assertEqual(
											    result.last_hidden_state.shape		,							(self.batch_size, expected_height, expected_width, expected_dim)			)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							A_			)->						List[str]:







											'''simple docstring'''






											UpperCamelCase					       =			DinatForImageClassification(A_			)
											model.to(A_			)
											model.eval()
											UpperCamelCase					       =			model(A_		,							labels=A_			)
											self.parent.assertEqual(result.logits.shape		,							(self.batch_size, self.num_labels)			)

											# test greyscale images
											UpperCamelCase					       =			1
											UpperCamelCase					       =			DinatForImageClassification(A_			)
											model.to(A_			)
											model.eval()

											UpperCamelCase					       =			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
											UpperCamelCase					       =			model(A_			)
											self.parent.assertEqual(result.logits.shape		,							(self.batch_size, self.num_labels)			)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							A_			)->						List[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			DinatBackbone(config=A_			)
											model.to(A_			)
											model.eval()
											UpperCamelCase					       =			model(A_			)

											# verify hidden states
											self.parent.assertEqual(len(result.feature_maps			)		,							len(config.out_features			)			)
											self.parent.assertListEqual(list(result.feature_maps[0].shape			)		,							[self.batch_size, model.channels[0], 16, 16]			)

											# verify channels
											self.parent.assertEqual(len(model.channels			)		,							len(config.out_features			)			)

											# verify backbone works with out_features=None
											UpperCamelCase					       =			None
											UpperCamelCase					       =			DinatBackbone(config=A_			)
											model.to(A_			)
											model.eval()
											UpperCamelCase					       =			model(A_			)

											# verify feature maps
											self.parent.assertEqual(len(result.feature_maps			)		,							1			)
											self.parent.assertListEqual(list(result.feature_maps[0].shape			)		,							[self.batch_size, model.channels[-1], 4, 4]			)

											# verify channels
											self.parent.assertEqual(len(model.channels			)		,							1			)







				def 		UpperCAmelCase_  (    self			)->						List[Any]:







											'''simple docstring'''






											UpperCamelCase					       =			self.prepare_config_and_inputs()
											UpperCamelCase       ,     UpperCamelCase       ,     UpperCamelCase					       =			config_and_inputs
											UpperCamelCase					       =			{'pixel_values': pixel_values}
											return config, inputs_dict


@require_natten
@require_torch
class        SCREAMING_SNAKE_CASE__	(			snake_case_				,    snake_case_				,    unittest.TestCase):
				lowerCAmelCase_			       =				(
				    (
				        DinatModel,
				        DinatForImageClassification,
				        DinatBackbone,
				    )
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase_			       =				(
				    {"""feature-extraction""": DinatModel, """image-classification""": DinatForImageClassification}
				    if is_torch_available()
				    else {}
				)
				lowerCAmelCase_			       =				False

				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False
				lowerCAmelCase_			       =				False







				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											UpperCamelCase					       =			DinatModelTester(self			)
											UpperCamelCase					       =			ConfigTester(self		,							config_class=A_		,							embed_dim=37			)







				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											self.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.create_and_test_config_with_num_labels()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()







				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											return







				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*A_			)







				def 		UpperCAmelCase_  (    self			)->						Any:







											'''simple docstring'''






											UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*A_			)







				def 		UpperCAmelCase_  (    self			)->						Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_backbone(*A_			)







				@unittest.skip(reason='Dinat does not use inputs_embeds'			)
				def 		UpperCAmelCase_  (    self			)->						Tuple:







											'''simple docstring'''






											pass







				@unittest.skip(reason='Dinat does not use feedforward chunking'			)
				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											pass







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







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																		UpperCamelCase					       =			model_class(A_			)
																		self.assertIsInstance(model.get_input_embeddings()		,							(nn.Module)			)
																		UpperCamelCase					       =			model.get_output_embeddings()
																		self.assertTrue(x is None or isinstance(A_		,							nn.Linear			)			)







				def 		UpperCAmelCase_  (    self			)->						int:







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()

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

																		UpperCamelCase					       =			['pixel_values']
																		self.assertListEqual(arg_names[:1]		,							A_			)







				def 		UpperCAmelCase_  (    self			)->						str:







											'''simple docstring'''






											self.skipTest('Dinat\'s attention operation is handled entirely by NATTEN.'			)







				def 		UpperCAmelCase_  (    self		,							A_		,							A_		,							A_		,							A_			)->						Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase					       =			model_class(A_			)
											model.to(A_			)
											model.eval()

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

											UpperCamelCase					       =			outputs.hidden_states

											UpperCamelCase					       =			getattr(
											    self.model_tester		,							'expected_num_hidden_layers'		,							len(self.model_tester.depths			) + 1			)
											self.assertEqual(len(A_			)		,							A_			)

											# Dinat has a different seq_length
											UpperCamelCase					       =			(
											    config.patch_size
											    if isinstance(config.patch_size		,							collections.abc.Iterable			)
											    else (config.patch_size, config.patch_size)
											)

											UpperCamelCase					       =			image_size[0] // patch_size[0]
											UpperCamelCase					       =			image_size[1] // patch_size[1]

											self.assertListEqual(
											    list(hidden_states[0].shape[-3:]			)		,							[height, width, self.model_tester.embed_dim]		,							)

											if model_class.__name__ != "DinatBackbone":
																		UpperCamelCase					       =			outputs.reshaped_hidden_states
																		self.assertEqual(len(A_			)		,							A_			)

																		UpperCamelCase       ,     UpperCamelCase       ,     UpperCamelCase       ,     UpperCamelCase					       =			reshaped_hidden_states[0].shape
																		UpperCamelCase					       =			(
																		    reshaped_hidden_states[0].view(A_		,							A_		,							A_		,							A_			).permute(0		,							2		,							3		,							1			)
																		)
																		self.assertListEqual(
																		    list(reshaped_hidden_states.shape[-3:]			)		,							[height, width, self.model_tester.embed_dim]		,							)







				def 		UpperCAmelCase_  (    self			)->						Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()

											UpperCamelCase					       =			(
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size		,							collections.abc.Iterable			)
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)

											for model_class in self.all_model_classes:
																		UpperCamelCase					       =			True
																		self.check_hidden_states_output(A_		,							A_		,							A_		,							A_			)

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

																		self.check_hidden_states_output(A_		,							A_		,							A_		,							A_			)







				@slow
				def 		UpperCAmelCase_  (    self			)->						Optional[int]:







											'''simple docstring'''






											for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCamelCase					       =			DinatModel.from_pretrained(A_			)
																		self.assertIsNotNone(A_			)







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







											'''simple docstring'''






											UpperCamelCase       ,     UpperCamelCase					       =			self.model_tester.prepare_config_and_inputs_for_common()

											UpperCamelCase					       =			_config_zero_init(A_			)
											for model_class in self.all_model_classes:
																		UpperCamelCase					       =			model_class(config=A_			)
																		for name, param in model.named_parameters():
																									if "embeddings" not in name and 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'''		,							)


@require_natten
@require_vision
@require_torch
class        SCREAMING_SNAKE_CASE__	(			unittest.TestCase):







				@cached_property
				def 		UpperCAmelCase_  (    self			)->						Union[str, Any]:







											'''simple docstring'''






											return AutoImageProcessor.from_pretrained('shi-labs/dinat-mini-in1k-224'			) if is_vision_available() else None







				@slow
				def 		UpperCAmelCase_  (    self			)->						Optional[int]:







											'''simple docstring'''






											UpperCamelCase					       =			DinatForImageClassification.from_pretrained('shi-labs/dinat-mini-in1k-224'			).to(A_			)
											UpperCamelCase					       =			self.default_image_processor

											UpperCamelCase					       =			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'			)
											UpperCamelCase					       =			image_processor(images=A_		,							return_tensors='pt'			).to(A_			)

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

											# verify the logits
											UpperCamelCase					       =			torch.Size((1, 1000)			)
											self.assertEqual(outputs.logits.shape		,							A_			)
											UpperCamelCase					       =			torch.tensor([-0.1_545, -0.7_667, 0.4_642]			).to(A_			)
											self.assertTrue(torch.allclose(outputs.logits[0, :3]		,							A_		,							atol=1e-4			)			)


@require_torch
@require_natten
class        SCREAMING_SNAKE_CASE__	(			unittest.TestCase				,    snake_case_):
				lowerCAmelCase_			       =				(DinatBackbone,) if is_torch_available() else ()
				lowerCAmelCase_			       =				DinatConfig







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







											'''simple docstring'''






											UpperCamelCase					       =			DinatModelTester(self			)






'''simple docstring'''
import argparse
import logging
import math
import time

import torch

from transformers import TransfoXLCorpus, TransfoXLLMHeadModel


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
lowerCAmelCase :  Optional[Any]      =      logging.getLogger(__name__)







def 						A_(			):
							UpperCamelCase					       =			argparse.ArgumentParser(description='PyTorch Transformer Language Model')
							parser.add_argument('--model_name'						,			type=A						,			default='transfo-xl-wt103'						,			help='pretrained model name')
							parser.add_argument(
							    '--split'						,			type=A						,			default='test'						,			choices=['all', 'valid', 'test']						,			help='which split to evaluate')
							parser.add_argument('--batch_size'						,			type=A						,			default=10						,			help='batch size')
							parser.add_argument('--tgt_len'						,			type=A						,			default=128						,			help='number of tokens to predict')
							parser.add_argument('--ext_len'						,			type=A						,			default=0						,			help='length of the extended context')
							parser.add_argument('--mem_len'						,			type=A						,			default=1600						,			help='length of the retained previous heads')
							parser.add_argument('--clamp_len'						,			type=A						,			default=1000						,			help='max positional embedding index')
							parser.add_argument('--no_cuda'						,			action='store_true'						,			help='Do not use CUDA even though CUA is available')
							parser.add_argument('--work_dir'						,			type=A						,			required=A						,			help='path to the work_dir')
							parser.add_argument('--no_log'						,			action='store_true'						,			help='do not log the eval result')
							parser.add_argument('--same_length'						,			action='store_true'						,			help='set same length attention with masking')
							parser.add_argument('--server_ip'						,			type=A						,			default=''						,			help='Can be used for distant debugging.')
							parser.add_argument('--server_port'						,			type=A						,			default=''						,			help='Can be used for distant debugging.')
							UpperCamelCase					       =			parser.parse_args()
							assert args.ext_len >= 0, "extended context length must be non-negative"

							if args.server_ip and args.server_port:
														# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
														import ptvsd

														print('Waiting for debugger attach')
														ptvsd.enable_attach(address=(args.server_ip, args.server_port)						,			redirect_output=A)
														ptvsd.wait_for_attach()

							UpperCamelCase					       =			torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu')
							logger.info('device: {}'.format(A))

							# Load a pre-processed dataset
							# You can also build the corpus yourself using TransfoXLCorpus methods
							# The pre-processing involve computing word frequencies to prepare the Adaptive input and SoftMax
							# and tokenizing the dataset
							# The pre-processed corpus is a convertion (using the conversion script )
							UpperCamelCase					       =			TransfoXLCorpus.from_pretrained(args.model_name)

							UpperCamelCase					       =			corpus.get_iterator('valid'						,			args.batch_size						,			args.tgt_len						,			device=A						,			ext_len=args.ext_len)
							UpperCamelCase					       =			corpus.get_iterator('test'						,			args.batch_size						,			args.tgt_len						,			device=A						,			ext_len=args.ext_len)

							# Load a pre-trained model
							UpperCamelCase					       =			TransfoXLLMHeadModel.from_pretrained(args.model_name)
							model.to(A)

							logger.info(
							    'Evaluating with bsz {} tgt_len {} ext_len {} mem_len {} clamp_len {}'.format(
							        args.batch_size						,			args.tgt_len						,			args.ext_len						,			args.mem_len						,			args.clamp_len))

							model.reset_memory_length(args.mem_len)
							if args.clamp_len > 0:
														UpperCamelCase					       =			args.clamp_len
							if args.same_length:
														UpperCamelCase					       =			True

							###############################################################################
							# Evaluation code
							###############################################################################
							def evaluate(A		:						Dict):
														# Turn on evaluation mode which disables dropout.
														model.eval()
														UpperCamelCase       ,     UpperCamelCase					       =			0, 0.0
														UpperCamelCase					       =			time.time()
														with torch.no_grad():
																					UpperCamelCase					       =			None
																					for idx, (data, target, seq_len) in enumerate(A):
																												UpperCamelCase					       =			model(A						,			lm_labels=A						,			mems=A)
																												UpperCamelCase       ,     UpperCamelCase       ,     UpperCamelCase					       =			ret
																												UpperCamelCase					       =			loss.mean()
																												total_loss += seq_len * loss.item()
																												total_len += seq_len
																					UpperCamelCase					       =			time.time() - start_time
														logger.info('Time : {:.2f}s, {:.2f}ms/segment'.format(A						,			1000 * total_time / (idx + 1)))
														return total_loss / total_len

							# Run on test data.
							if args.split == "all":
														UpperCamelCase					       =			evaluate(A)
														UpperCamelCase					       =			evaluate(A)
							elif args.split == "valid":
														UpperCamelCase					       =			evaluate(A)
														UpperCamelCase					       =			None
							elif args.split == "test":
														UpperCamelCase					       =			evaluate(A)
														UpperCamelCase					       =			None

							def format_log(A		:						Any						,			A		:						Optional[Any]):
														UpperCamelCase					       =			'| {0} loss {1:5.2f} | {0} ppl {2:9.3f} '.format(A						,			A						,			math.exp(A))
														return log_str

							UpperCamelCase					       =			''
							if valid_loss is not None:
														log_str += format_log(A						,			'valid')
							if test_loss is not None:
														log_str += format_log(A						,			'test')

							logger.info('=' * 100)
							logger.info(A)
							logger.info('=' * 100)


if __name__ == "__main__":
						main()




