
import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling




def 							a		(				snake_case__:   Dict   ):


				'''simple docstring'''
				lowercase_	    =   [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
				lowercase_	    =   True if '''large''' in model_name or '''huge''' in model_name else False
				lowercase_	    =   True if '''large''' in model_name or '''huge''' in model_name else False
				lowercase_	    =   True if '''large''' in model_name or '''huge''' in model_name else False

				if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
								if "fl3" in model_name:
												lowercase_	    =   [3, 3, 3, 3]
												lowercase_	    =   [5, 5, 5, 5]
								elif "fl4" in model_name:
												lowercase_	    =   [4, 4, 4, 4]
												lowercase_	    =   [3, 3, 3, 3]

				if "tiny" in model_name or "small" in model_name or "base" in model_name:
								lowercase_	    =   [3, 3, 3, 3]
								if "lrf" in model_name:
												lowercase_	    =   [3, 3, 3, 3]
								else:
												lowercase_	    =   [2, 2, 2, 2]

				if "tiny" in model_name:
								lowercase_	    =   96
				elif "small" in model_name:
								lowercase_	    =   96
				elif "base" in model_name:
								lowercase_	    =   128
				elif "large" in model_name:
								lowercase_	    =   192
				elif "xlarge" in model_name:
								lowercase_	    =   256
				elif "huge" in model_name:
								lowercase_	    =   352

				# set label information
				lowercase_	    =   '''huggingface/label-files'''
				if "large" in model_name or "huge" in model_name:
								lowercase_	    =   '''imagenet-22k-id2label.json'''
				else:
								lowercase_	    =   '''imagenet-1k-id2label.json'''

				lowercase_	    =   json.load(open(hf_hub_download(snake_case__		,  snake_case__		,  repo_type='''dataset'''   )		,  '''r'''   )   )
				lowercase_	    =   {int(snake_case__   ): v for k, v in idalabel.items()}
				lowercase_	    =   {v: k for k, v in idalabel.items()}

				lowercase_	    =   FocalNetConfig(
				    embed_dim=snake_case__		,  depths=snake_case__		,  focal_levels=snake_case__		,  focal_windows=snake_case__		,  use_conv_embed=snake_case__		,  idalabel=snake_case__		,  labelaid=snake_case__		,  use_post_layernorm=snake_case__		,  use_layerscale=snake_case__		,  )

				return config




def 							a		(				snake_case__:   Any   ):


				'''simple docstring'''
				if "patch_embed.proj" in name:
								lowercase_	    =   name.replace('''patch_embed.proj'''		,  '''embeddings.patch_embeddings.projection'''   )
				if "patch_embed.norm" in name:
								lowercase_	    =   name.replace('''patch_embed.norm'''		,  '''embeddings.norm'''   )
				if "layers" in name:
								lowercase_	    =   '''encoder.''' + name
				if "encoder.layers" in name:
								lowercase_	    =   name.replace('''encoder.layers'''		,  '''encoder.stages'''   )
				if "downsample.proj" in name:
								lowercase_	    =   name.replace('''downsample.proj'''		,  '''downsample.projection'''   )
				if "blocks" in name:
								lowercase_	    =   name.replace('''blocks'''		,  '''layers'''   )
				if "modulation.f.weight" in name or "modulation.f.bias" in name:
								lowercase_	    =   name.replace('''modulation.f'''		,  '''modulation.projection_in'''   )
				if "modulation.h.weight" in name or "modulation.h.bias" in name:
								lowercase_	    =   name.replace('''modulation.h'''		,  '''modulation.projection_context'''   )
				if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
								lowercase_	    =   name.replace('''modulation.proj'''		,  '''modulation.projection_out'''   )

				if name == "norm.weight":
								lowercase_	    =   '''layernorm.weight'''
				if name == "norm.bias":
								lowercase_	    =   '''layernorm.bias'''

				if "head" in name:
								lowercase_	    =   name.replace('''head'''		,  '''classifier'''   )
				else:
								lowercase_	    =   '''focalnet.''' + name

				return name




def 							a		(				snake_case__:   str		,  snake_case__:   str		,  snake_case__:   Any=False   ):


				'''simple docstring'''
				# fmt: off
				lowercase_	    =   {
				    '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
				    '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
				    '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
				    '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
				    '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
				    '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
				    '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
				    '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
				    '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
				    '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
				}
				# fmt: on

				lowercase_	    =   model_name_to_url[model_name]
				print('''Checkpoint URL: '''		,  snake_case__   )
				lowercase_	    =   torch.hub.load_state_dict_from_url(snake_case__		,  map_location='''cpu'''   )['''model''']

				# rename keys
				for key in state_dict.copy().keys():
								lowercase_	    =   state_dict.pop(snake_case__   )
								lowercase_	    =   val

				lowercase_	    =   get_focalnet_config(snake_case__   )
				lowercase_	    =   FocalNetForImageClassification(snake_case__   )
				model.eval()

				# load state dict
				model.load_state_dict(snake_case__   )

				# verify conversion
				lowercase_	    =   '''http://images.cocodataset.org/val2017/000000039769.jpg'''

				lowercase_	    =   BitImageProcessor(
				    do_resize=snake_case__		,  size={'''shortest_edge''': 256}		,  resample=PILImageResampling.BILINEAR		,  do_center_crop=snake_case__		,  crop_size=224		,  do_normalize=snake_case__		,  image_mean=snake_case__		,  image_std=snake_case__		,  )
				lowercase_	    =   Image.open(requests.get(snake_case__		,  stream=snake_case__   ).raw   )
				lowercase_	    =   processor(images=snake_case__		,  return_tensors='''pt'''   )

				lowercase_	    =   transforms.Compose(
				    [
				        transforms.Resize(256   ),
				        transforms.CenterCrop(224   ),
				        transforms.ToTensor(),
				        transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]		,  std=[0.2_2_9, 0.2_2_4, 0.2_2_5]   ),
				    ]   )

				lowercase_	    =   image_transforms(snake_case__   ).unsqueeze(0   )

				# verify pixel_values
				assert torch.allclose(inputs.pixel_values		,  snake_case__		,  atol=1e-4   )

				lowercase_	    =   model(**snake_case__   )

				lowercase_	    =   outputs.logits.argmax(-1   ).item()
				print('''Predicted class:'''		,  model.config.idalabel[predicted_class_idx]   )

				print('''First values of logits:'''		,  outputs.logits[0, :3]   )

				if model_name == "focalnet-tiny":
								lowercase_	    =   torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]   )
				elif model_name == "focalnet-tiny-lrf":
								lowercase_	    =   torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]   )
				elif model_name == "focalnet-small":
								lowercase_	    =   torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]   )
				elif model_name == "focalnet-small-lrf":
								lowercase_	    =   torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]   )
				elif model_name == "focalnet-base":
								lowercase_	    =   torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]   )
				elif model_name == "focalnet-base-lrf":
								lowercase_	    =   torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]   )
				assert torch.allclose(outputs.logits[0, :3]		,  snake_case__		,  atol=1e-4   )
				print('''Looks ok!'''   )

				if pytorch_dump_folder_path is not None:
								print(F'''Saving model and processor of {model_name} to {pytorch_dump_folder_path}'''   )
								model.save_pretrained(snake_case__   )
								processor.save_pretrained(snake_case__   )

				if push_to_hub:
								print(F'''Pushing model and processor of {model_name} to the hub...'''   )
								model.push_to_hub(F'''{model_name}'''   )
								processor.push_to_hub(F'''{model_name}'''   )


if __name__ == "__main__":
						__a         =   argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--model_name',
						    default='focalnet-tiny',
						    type=str,
						    help='Name of the FocalNet model you\'d like to convert.',
						)
						parser.add_argument(
						    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
						)
						parser.add_argument(
						    '--push_to_hub',
						    action='store_true',
						    help='Whether to push the model and processor to the hub.',
						)

						__a         =   parser.parse_args()
						convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
				import torch

				from transformers import (
				    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    SqueezeBertForMaskedLM,
				    SqueezeBertForMultipleChoice,
				    SqueezeBertForQuestionAnswering,
				    SqueezeBertForSequenceClassification,
				    SqueezeBertForTokenClassification,
				    SqueezeBertModel,
				)



class        lowerCamelCase__				(   UpperCAmelCase				):





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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size						)

										_lowerCAmelCase							=							None
										if self.use_input_mask:
																	_lowerCAmelCase							=							random_attention_mask([self.batch_size, self.seq_length]						)

										_lowerCAmelCase							=							None
										_lowerCAmelCase							=							None
										_lowerCAmelCase							=							None
										if self.use_labels:
																	_lowerCAmelCase							=							ids_tensor([self.batch_size]  ,  self.type_sequence_label_size						)
																	_lowerCAmelCase							=							ids_tensor([self.batch_size, self.seq_length]  ,  self.num_labels						)
																	_lowerCAmelCase							=							ids_tensor([self.batch_size]  ,  self.num_choices						)

										_lowerCAmelCase							=							self.get_config()

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										return SqueezeBertConfig(
										    embedding_size=self.hidden_size  ,  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  ,  attention_probs_dropout_prob=self.hidden_dropout_prob  ,  attention_dropout=self.attention_probs_dropout_prob  ,  max_position_embeddings=self.max_position_embeddings  ,  initializer_range=self.initializer_range  ,  q_groups=self.q_groups  ,  k_groups=self.k_groups  ,  v_groups=self.v_groups  ,  post_attention_groups=self.post_attention_groups  ,  intermediate_groups=self.intermediate_groups  ,  output_groups=self.output_groups  ,  )





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Dict  ,  lowercase__	:	Any  ,  lowercase__	:	str  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Any  ,  lowercase__	:	Tuple						):
										_lowerCAmelCase							=							SqueezeBertModel(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  lowercase__						)
										_lowerCAmelCase							=							model(lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]						):
										_lowerCAmelCase							=							SqueezeBertForMaskedLM(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.vocab_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Tuple  ,  lowercase__	:	Dict  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	int						):
										_lowerCAmelCase							=							SqueezeBertForQuestionAnswering(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(
										    lowercase__  ,  attention_mask=lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__						)
										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   SCREAMING_SNAKE_CASE__     (       self	:	Dict  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Tuple						):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							SqueezeBertForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.num_labels)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	Tuple  ,  lowercase__	:	Any  ,  lowercase__	:	Optional[int]						):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							SqueezeBertForTokenClassification(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__  ,  attention_mask=lowercase__  ,  labels=lowercase__						)
										self.parent.assertEqual(result.logits.shape  ,  (self.batch_size, self.seq_length, self.num_labels)						)





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





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							self.prepare_config_and_inputs()
										((_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase)			,		(_lowerCAmelCase))							=							config_and_inputs
										_lowerCAmelCase							=							{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict







@require_torch
class        lowerCamelCase__				(   UpperCAmelCase       ,UpperCAmelCase       ,unittest.TestCase				):
			UpperCamelCase__    							=(
			    (
			        SqueezeBertModel,
			        SqueezeBertForMaskedLM,
			        SqueezeBertForMultipleChoice,
			        SqueezeBertForQuestionAnswering,
			        SqueezeBertForSequenceClassification,
			        SqueezeBertForTokenClassification,
			    )
			    if is_torch_available()
			    else None
			)
			UpperCamelCase__    							=(
			    {
			        "feature-extraction": SqueezeBertModel,
			        "fill-mask": SqueezeBertForMaskedLM,
			        "question-answering": SqueezeBertForQuestionAnswering,
			        "text-classification": SqueezeBertForSequenceClassification,
			        "token-classification": SqueezeBertForTokenClassification,
			        "zero-shot": SqueezeBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			UpperCamelCase__    							=False
			UpperCamelCase__    							=True
			UpperCamelCase__    							=False





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							SqueezeBertModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  dim=37						)





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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Tuple						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_masked_lm(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_question_answering(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_sequence_classification(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_token_classification(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_squeezebert_for_multiple_choice(*lowercase__						)





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_lowerCAmelCase							=							SqueezeBertModel.from_pretrained(lowercase__						)
																	self.assertIsNotNone(lowercase__						)







@require_sentencepiece
@require_tokenizers
@require_torch
class        lowerCamelCase__				(   unittest.TestCase				):





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_lowerCAmelCase							=							SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli'						)

										_lowerCAmelCase							=							torch.tensor([[1, 2_94_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 13, 15_88, 2]]						)
										_lowerCAmelCase							=							model(lowercase__						)[0]
										_lowerCAmelCase							=							torch.Size((1, 3)						)
										self.assertEqual(output.shape  ,  lowercase__						)
										_lowerCAmelCase							=							torch.tensor([[0.6_4_0_1, -0.0_3_4_9, -0.6_0_4_1]]						)
										self.assertTrue(torch.allclose(lowercase__  ,  lowercase__  ,  atol=1e-4						)						)






