



"""simple docstring"""
import inspect
import tempfile
import unittest

from huggingface_hub import hf_hub_download

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

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


A:   Optional[Any]								=		1E-4

if is_torch_available():
      import torch

      from transformers import AutoformerConfig, AutoformerForPrediction, AutoformerModel
      from transformers.models.autoformer.modeling_autoformer import AutoformerDecoder, AutoformerEncoder



@require_torch
class       SCREAMING_SNAKE_CASE__  :

   def __init__(   self      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE=16      ,     _SCREAMING_SNAKE_CASE=13      ,     _SCREAMING_SNAKE_CASE=7      ,     _SCREAMING_SNAKE_CASE=14      ,     _SCREAMING_SNAKE_CASE=10      ,     _SCREAMING_SNAKE_CASE=19      ,     _SCREAMING_SNAKE_CASE=5      ,     _SCREAMING_SNAKE_CASE=4      ,     _SCREAMING_SNAKE_CASE=True      ,     _SCREAMING_SNAKE_CASE=16      ,     _SCREAMING_SNAKE_CASE=2      ,     _SCREAMING_SNAKE_CASE=4      ,     _SCREAMING_SNAKE_CASE=4      ,     _SCREAMING_SNAKE_CASE="gelu"      ,     _SCREAMING_SNAKE_CASE=0.1      ,     _SCREAMING_SNAKE_CASE=0.1      ,     _SCREAMING_SNAKE_CASE=[1, 2, 3, 4, 5]      ,     _SCREAMING_SNAKE_CASE=25      ,     _SCREAMING_SNAKE_CASE=5      ,     )  ->      List[str]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[Any]	=	d_model
    UpperCAmelCase	:					Optional[int]	=	parent
    UpperCAmelCase	:					Tuple	=	batch_size
    UpperCAmelCase	:					Optional[int]	=	prediction_length
    UpperCAmelCase	:					Optional[Any]	=	context_length
    UpperCAmelCase	:					Optional[int]	=	cardinality
    UpperCAmelCase	:					Union[str, Any]	=	num_time_features
    UpperCAmelCase	:					Union[str, Any]	=	lags_sequence
    UpperCAmelCase	:					Dict	=	embedding_dimension
    UpperCAmelCase	:					Dict	=	is_training
    UpperCAmelCase	:					List[Any]	=	hidden_size
    UpperCAmelCase	:					List[str]	=	num_hidden_layers
    UpperCAmelCase	:					Any	=	num_attention_heads
    UpperCAmelCase	:					List[Any]	=	intermediate_size
    UpperCAmelCase	:					Optional[Any]	=	hidden_act
    UpperCAmelCase	:					str	=	hidden_dropout_prob
    UpperCAmelCase	:					str	=	attention_probs_dropout_prob

    UpperCAmelCase	:					Union[str, Any]	=	context_length
    UpperCAmelCase	:					int	=	prediction_length + label_length
    UpperCAmelCase	:					Any	=	label_length

    UpperCAmelCase	:					str	=	moving_average
    UpperCAmelCase	:					int	=	autocorrelation_factor

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





    '''simple docstring'''




    return AutoformerConfig(
        d_model=self.d_model      ,     encoder_layers=self.num_hidden_layers      ,     decoder_layers=self.num_hidden_layers      ,     encoder_attention_heads=self.num_attention_heads      ,     decoder_attention_heads=self.num_attention_heads      ,     encoder_ffn_dim=self.intermediate_size      ,     decoder_ffn_dim=self.intermediate_size      ,     dropout=self.hidden_dropout_prob      ,     attention_dropout=self.attention_probs_dropout_prob      ,     prediction_length=self.prediction_length      ,     context_length=self.context_length      ,     label_length=self.label_length      ,     lags_sequence=self.lags_sequence      ,     num_time_features=self.num_time_features      ,     num_static_categorical_features=1      ,     cardinality=[self.cardinality]      ,     embedding_dimension=[self.embedding_dimension]      ,     moving_average=self.moving_average      ,     )

   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE						)  ->      List[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	config.context_length + max(config.lags_sequence						)

    UpperCAmelCase	:					Tuple	=	ids_tensor([self.batch_size, 1]      ,     config.cardinality[0]						)
    UpperCAmelCase	:					Optional[Any]	=	floats_tensor([self.batch_size, _past_length, config.num_time_features]						)
    UpperCAmelCase	:					str	=	floats_tensor([self.batch_size, _past_length]						)
    UpperCAmelCase	:					Optional[Any]	=	floats_tensor([self.batch_size, _past_length]						) > 0.5

    # decoder inputs
    UpperCAmelCase	:					str	=	floats_tensor([self.batch_size, config.prediction_length, config.num_time_features]						)
    UpperCAmelCase	:					Dict	=	floats_tensor([self.batch_size, config.prediction_length]						)

    UpperCAmelCase	:					Dict	=	{
        """past_values""": past_values,
        """static_categorical_features""": static_categorical_features,
        """past_time_features""": past_time_features,
        """past_observed_mask""": past_observed_mask,
        """future_time_features""": future_time_features,
        """future_values""": future_values,
    }
    return inputs_dict

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Tuple	=	self.get_config()
    UpperCAmelCase	:					List[str]	=	self.prepare_autoformer_inputs_dict(_SCREAMING_SNAKE_CASE						)
    return config, inputs_dict

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





    '''simple docstring'''




    UpperCAmelCase	,		UpperCAmelCase	:					int	=	self.prepare_config_and_inputs()
    return config, inputs_dict




   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)  ->      List[str]:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	AutoformerModel(config=_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						).eval()
    UpperCAmelCase	:					List[str]	=	model(**_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					List[str]	=	outputs.encoder_last_hidden_state
    UpperCAmelCase	:					Optional[Any]	=	outputs.last_hidden_state

    with tempfile.TemporaryDirectory() as tmpdirname:
     UpperCAmelCase	:					Any	=	model.get_encoder()
     encoder.save_pretrained(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Optional[int]	=	AutoformerEncoder.from_pretrained(_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	,		UpperCAmelCase	:					Union[str, Any]	=	model.create_network_inputs(**_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	,		UpperCAmelCase	:					Tuple	=	model.decomposition_layer(transformer_inputs[:, : config.context_length, ...]						)

    UpperCAmelCase	:					int	=	torch.cat(
        (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])      ,     dim=-1      ,     )
    UpperCAmelCase	:					Optional[Any]	=	encoder(inputs_embeds=_SCREAMING_SNAKE_CASE						)[0]
    self.parent.assertTrue((encoder_last_hidden_state_a - encoder_last_hidden_state).abs().max().item() < 1E-3						)

    UpperCAmelCase	:					str	=	(
        torch.mean(transformer_inputs[:, : config.context_length, ...]      ,     dim=1						)
        .unsqueeze(1						)
        .repeat(1      ,     config.prediction_length      ,     1						)
    )
    UpperCAmelCase	:					Optional[int]	=	torch.zeros(
        [transformer_inputs.shape[0], config.prediction_length, transformer_inputs.shape[2]]      ,     device=enc_input.device      ,     )

    UpperCAmelCase	:					str	=	torch.cat(
        (
            torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)      ,     dim=1						),
            feature[:, config.context_length - config.label_length :, ...],
        )      ,     dim=-1      ,     )
    UpperCAmelCase	:					Optional[int]	=	torch.cat(
        (
            torch.cat((trend_input[:, -config.label_length :, ...], mean)      ,     dim=1						),
            feature[:, config.context_length - config.label_length :, ...],
        )      ,     dim=-1      ,     )

    with tempfile.TemporaryDirectory() as tmpdirname:
     UpperCAmelCase	:					str	=	model.get_decoder()
     decoder.save_pretrained(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Optional[Any]	=	AutoformerDecoder.from_pretrained(_SCREAMING_SNAKE_CASE						).to(_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Dict	=	decoder(
        trend=_SCREAMING_SNAKE_CASE      ,     inputs_embeds=_SCREAMING_SNAKE_CASE      ,     encoder_hidden_states=_SCREAMING_SNAKE_CASE      ,     )[0]

    self.parent.assertTrue((last_hidden_state_a - last_hidden_state).abs().max().item() < 1E-3						)



@require_torch
class       SCREAMING_SNAKE_CASE__  (					UpperCAmelCase__      ,			UpperCAmelCase__      ,			unittest.TestCase ):
   __lowerCAmelCase  :   Any						=						(AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
   __lowerCAmelCase  :   Any						=						(AutoformerForPrediction,) if is_torch_available() else ()
   __lowerCAmelCase  :   int						=						{'feature-extraction': AutoformerModel} if is_torch_available() else {}
   __lowerCAmelCase  :   List[Any]						=						False
   __lowerCAmelCase  :   Tuple						=						False
   __lowerCAmelCase  :   Any						=						False
   __lowerCAmelCase  :   Union[str, Any]						=						False
   __lowerCAmelCase  :   Optional[Any]						=						False
   __lowerCAmelCase  :   Union[str, Any]						=						False

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[int]:





    '''simple docstring'''




    UpperCAmelCase	:					int	=	AutoformerModelTester(self						)
    UpperCAmelCase	:					Tuple	=	ConfigTester(self      ,     config_class=_SCREAMING_SNAKE_CASE      ,     has_text_modality=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    self.config_tester.run_common_tests()

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	,		UpperCAmelCase	:					Tuple	=	self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
     UpperCAmelCase	:					int	=	model_class(_SCREAMING_SNAKE_CASE						)

     with tempfile.TemporaryDirectory() as tmpdirname:
      model.save_pretrained(_SCREAMING_SNAKE_CASE						)
      UpperCAmelCase	,		UpperCAmelCase	:					Union[str, Any]	=	model_class.from_pretrained(_SCREAMING_SNAKE_CASE      ,     output_loading_info=_SCREAMING_SNAKE_CASE						)
     self.assertEqual(info["""missing_keys"""]      ,     []						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					List[Any]	=	self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_encoder_decoder_model_standalone(*_SCREAMING_SNAKE_CASE						)

   @unittest.skip(reason="""Model has no tokens embeddings"""						)
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Union[str, Any]:





    '''simple docstring'''




    pass

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					Union[str, Any]	=	inspect.signature(getattr(_SCREAMING_SNAKE_CASE      ,     """forward"""						)						)
    # The main input is the name of the argument after `self`
    UpperCAmelCase	:					Any	=	list(model_signature.parameters.keys()						)[1]
    self.assertEqual(AutoformerModel.main_input_name      ,     _SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	,		UpperCAmelCase	:					str	=	self.model_tester.prepare_config_and_inputs_for_common()

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

     UpperCAmelCase	:					Optional[int]	=	[
         """past_values""",
         """past_time_features""",
         """past_observed_mask""",
         """static_categorical_features""",
         """static_real_features""",
         """future_values""",
         """future_time_features""",
     ]

     if model.__class__.__name__ in ["AutoformerForPrediction"]:
      expected_arg_names.append("""future_observed_mask"""						)

     expected_arg_names.extend(
         [
             """decoder_attention_mask""",
             """head_mask""",
             """decoder_head_mask""",
             """cross_attn_head_mask""",
             """encoder_outputs""",
             """past_key_values""",
             """output_hidden_states""",
             """output_attentions""",
             """use_cache""",
             """return_dict""",
         ]						)

     self.assertListEqual(arg_names[: len(_SCREAMING_SNAKE_CASE						)]      ,     _SCREAMING_SNAKE_CASE						)

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





    '''simple docstring'''




    UpperCAmelCase	,		UpperCAmelCase	:					Any	=	self.model_tester.prepare_config_and_inputs_for_common()
    UpperCAmelCase	:					List[Any]	=	True

    UpperCAmelCase	:					List[str]	=	getattr(self.model_tester      ,     """seq_length"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Any	=	getattr(self.model_tester      ,     """decoder_seq_length"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Union[str, Any]	=	getattr(self.model_tester      ,     """encoder_seq_length"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Tuple	=	getattr(self.model_tester      ,     """d_model"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Optional[Any]	=	getattr(self.model_tester      ,     """num_attention_heads"""      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Optional[int]	=	d_model // num_attention_heads

    for model_class in self.all_model_classes:
     UpperCAmelCase	:					Any	=	True
     UpperCAmelCase	:					List[Any]	=	False
     UpperCAmelCase	:					Union[str, Any]	=	True
     UpperCAmelCase	:					List[Any]	=	model_class(_SCREAMING_SNAKE_CASE						)
     model.to(_SCREAMING_SNAKE_CASE						)
     model.eval()
     with torch.no_grad():
      UpperCAmelCase	:					List[str]	=	model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)
     UpperCAmelCase	:					Optional[int]	=	outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)

     # check that output_attentions also work using config
     del inputs_dict["output_attentions"]
     UpperCAmelCase	:					str	=	True
     UpperCAmelCase	:					Tuple	=	model_class(_SCREAMING_SNAKE_CASE						)
     model.to(_SCREAMING_SNAKE_CASE						)
     model.eval()
     with torch.no_grad():
      UpperCAmelCase	:					Optional[int]	=	model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)
     UpperCAmelCase	:					Any	=	outputs.encoder_attentions
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)

     self.assertListEqual(
         list(attentions[0].shape[-3:]						)      ,     [self.model_tester.num_attention_heads, encoder_seq_length, dim]      ,     )
     UpperCAmelCase	:					Optional[Any]	=	len(_SCREAMING_SNAKE_CASE						)

     UpperCAmelCase	:					int	=	7

     if "last_hidden_state" in outputs:
      correct_outlen += 1

     if "trend" in outputs:
      correct_outlen += 1

     if "past_key_values" in outputs:
      correct_outlen += 1  # past_key_values have been returned

     if "loss" in outputs:
      correct_outlen += 1

     if "params" in outputs:
      correct_outlen += 1

     self.assertEqual(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)

     # decoder attentions
     UpperCAmelCase	:					str	=	outputs.decoder_attentions
     self.assertIsInstance(_SCREAMING_SNAKE_CASE      ,     (list, tuple)						)
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)
     self.assertListEqual(
         list(decoder_attentions[0].shape[-3:]						)      ,     [self.model_tester.num_attention_heads, decoder_seq_length, dim]      ,     )

     # cross attentions
     UpperCAmelCase	:					Dict	=	outputs.cross_attentions
     self.assertIsInstance(_SCREAMING_SNAKE_CASE      ,     (list, tuple)						)
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)
     self.assertListEqual(
         list(cross_attentions[0].shape[-3:]						)      ,     [self.model_tester.num_attention_heads, decoder_seq_length, dim]      ,     )

    # Check attention is always last and order is fine
    UpperCAmelCase	:					int	=	True
    UpperCAmelCase	:					Optional[Any]	=	True
    UpperCAmelCase	:					List[Any]	=	model_class(_SCREAMING_SNAKE_CASE						)
    model.to(_SCREAMING_SNAKE_CASE						)
    model.eval()
    with torch.no_grad():
     UpperCAmelCase	:					Dict	=	model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)

    self.assertEqual(out_len + 2      ,     len(_SCREAMING_SNAKE_CASE						)						)

    UpperCAmelCase	:					Tuple	=	outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions

    self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)
    self.assertListEqual(
        list(self_attentions[0].shape[-3:]						)      ,     [self.model_tester.num_attention_heads, encoder_seq_length, dim]      ,     )




   @is_flaky()
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    super().test_retain_grad_hidden_states_attentions()

def       _snake_case						(						UpperCamelCase						:		Any="train-batch.pt"						):
 UpperCAmelCase	:					List[str]	=	hf_hub_download(repo_id="""hf-internal-testing/tourism-monthly-batch"""			,		filename=UpperCamelCase			,		repo_type="""dataset"""						)
 UpperCAmelCase	:					List[str]	=	torch.load(UpperCamelCase			,		map_location=UpperCamelCase						)
 return batch



@require_torch
@slow
class       SCREAMING_SNAKE_CASE__  (					unittest.TestCase ):

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





    '''simple docstring'''




    UpperCAmelCase	:					str	=	AutoformerModel.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Union[str, Any]	=	prepare_batch()

    with torch.no_grad():
     UpperCAmelCase	:					Any	=	model(
         past_values=batch["""past_values"""]      ,     past_time_features=batch["""past_time_features"""]      ,     past_observed_mask=batch["""past_observed_mask"""]      ,     static_categorical_features=batch["""static_categorical_features"""]      ,     future_values=batch["""future_values"""]      ,     future_time_features=batch["""future_time_features"""]      ,     )[0]

    UpperCAmelCase	:					Union[str, Any]	=	torch.Size(
        (64, model.config.prediction_length + model.config.label_length, model.config.feature_size)						)
    self.assertEqual(output.shape      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					str	=	torch.tensor(
        [[0.3593, -1.3398, 0.6330], [0.2279, 1.5396, -0.1792], [0.0450, 1.3225, -0.2335]]      ,     device=_SCREAMING_SNAKE_CASE						)
    self.assertTrue(torch.allclose(output[0, :3, :3]      ,     _SCREAMING_SNAKE_CASE      ,     atol=_SCREAMING_SNAKE_CASE						)						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Any	=	AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Dict	=	prepare_batch("""val-batch.pt"""						)
    with torch.no_grad():
     UpperCAmelCase	:					Dict	=	model(
         past_values=batch["""past_values"""]      ,     past_time_features=batch["""past_time_features"""]      ,     past_observed_mask=batch["""past_observed_mask"""]      ,     static_categorical_features=batch["""static_categorical_features"""]      ,     ).encoder_last_hidden_state
    UpperCAmelCase	:					List[str]	=	torch.Size((64, model.config.context_length, model.config.d_model)						)
    self.assertEqual(output.shape      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Union[str, Any]	=	torch.tensor(
        [[-0.0734, -0.9036, 0.8358], [4.7186, 2.4113, 1.9581], [1.7953, 2.3558, 1.2970]]      ,     device=_SCREAMING_SNAKE_CASE						)
    self.assertTrue(torch.allclose(output[0, :3, :3]      ,     _SCREAMING_SNAKE_CASE      ,     atol=_SCREAMING_SNAKE_CASE						)						)




   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Tuple:





    '''simple docstring'''




    UpperCAmelCase	:					str	=	AutoformerForPrediction.from_pretrained("""huggingface/autoformer-tourism-monthly"""						).to(_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Optional[Any]	=	prepare_batch("""val-batch.pt"""						)
    with torch.no_grad():
     UpperCAmelCase	:					Union[str, Any]	=	model.generate(
         static_categorical_features=batch["""static_categorical_features"""]      ,     past_time_features=batch["""past_time_features"""]      ,     past_values=batch["""past_values"""]      ,     future_time_features=batch["""future_time_features"""]      ,     past_observed_mask=batch["""past_observed_mask"""]      ,     )
    UpperCAmelCase	:					List[Any]	=	torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length)						)
    self.assertEqual(outputs.sequences.shape      ,     _SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	:					Tuple	=	torch.tensor([3130.6763, 4056.5293, 7053.0786]      ,     device=_SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Dict	=	outputs.sequences.mean(dim=1						)
    self.assertTrue(torch.allclose(mean_prediction[0, -3:]      ,     _SCREAMING_SNAKE_CASE      ,     rtol=1E-1						)						)





"""simple docstring"""
import argparse
import json
from pathlib import Path

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

from transformers import (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
A:   Optional[int]								=		logging.get_logger(__name__)

A:   Optional[int]								=		torch.device("cpu")

def       _snake_case						(						):
 UpperCAmelCase	:					List[Any]	=	"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 UpperCAmelCase	:					Tuple	=	Image.open(requests.get(UpperCamelCase			,		stream=UpperCamelCase						).raw						)
 return im

def       _snake_case						(						UpperCamelCase						:		int						):
 if swiftformer_name == "swiftformer_xs":
  return torch.tensor([-2.1_7_0_3e0_0, 2.1_1_0_7e0_0, -2.0_8_1_1e0_0, 8.8_6_8_5e-0_1, 2.4_3_6_0e-0_1]						)

 elif swiftformer_name == "swiftformer_s":
  return torch.tensor([3.9_6_3_6e-0_1, 2.3_4_7_8e-0_1, -1.6_9_6_3e0_0, -1.7_3_8_1e0_0, -8.6_3_3_7e-0_1]						)

 elif swiftformer_name == "swiftformer_l1":
  return torch.tensor([-4.2_7_6_8e-0_1, -4.7_4_2_9e-0_1, -1.0_8_9_7e0_0, -1.0_2_4_8e0_0, 3.5_5_2_3e-0_2]						)

 elif swiftformer_name == "swiftformer_l3":
  return torch.tensor([-2.5_3_3_0e-0_1, 2.4_2_1_1e-0_1, -6.0_1_8_5e-0_1, -8.2_7_8_9e-0_1, -6.0_4_4_6e-0_2]						)

def       _snake_case						(						UpperCamelCase						:		Optional[Any]			,		UpperCamelCase						:		List[str]			,		UpperCamelCase						:		str						):
 UpperCAmelCase	:					int	=	dct.pop(UpperCamelCase						)
 UpperCAmelCase	:					Any	=	val

def       _snake_case						(						UpperCamelCase						:		Union[str, Any]						):
 UpperCAmelCase	:					Optional[int]	=	[]
 for k in state_dict.keys():
  UpperCAmelCase	:					Optional[Any]	=	k
  if ".pwconv" in k:
   UpperCAmelCase	:					int	=	k_new.replace(""".pwconv"""			,		""".point_wise_conv"""						)
  if ".dwconv" in k:
   UpperCAmelCase	:					Tuple	=	k_new.replace(""".dwconv"""			,		""".depth_wise_conv"""						)
  if ".Proj." in k:
   UpperCAmelCase	:					List[Any]	=	k_new.replace(""".Proj."""			,		""".proj."""						)
  if "patch_embed" in k_new:
   UpperCAmelCase	:					Any	=	k_new.replace("""patch_embed"""			,		"""swiftformer.patch_embed.patch_embedding"""						)
  if "network" in k_new:
   UpperCAmelCase	:					int	=	k_new.split("""."""						)
   if ls[2].isdigit():
    UpperCAmelCase	:					List[Any]	=	"""swiftformer.encoder.network.""" + ls[1] + """.blocks.""" + ls[2] + """.""" + """.""".join(ls[3:]						)
   else:
    UpperCAmelCase	:					Any	=	k_new.replace("""network"""			,		"""swiftformer.encoder.network"""						)
  rename_keys.append((k, k_new)						)
 return rename_keys

@torch.no_grad()
def       _snake_case						(						UpperCamelCase						:		int			,		UpperCamelCase						:		int			,		UpperCamelCase						:		Optional[int]						):
 UpperCAmelCase	:					List[Any]	=	SwiftFormerConfig()

 # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
 UpperCAmelCase	:					Optional[Any]	=	1000
 UpperCAmelCase	:					Tuple	=	"""huggingface/label-files"""
 UpperCAmelCase	:					List[str]	=	"""imagenet-1k-id2label.json"""
 UpperCAmelCase	:					Dict	=	json.load(open(hf_hub_download(UpperCamelCase			,		UpperCamelCase			,		repo_type="""dataset"""						)			,		"""r"""						)						)
 UpperCAmelCase	:					Dict	=	{int(UpperCamelCase						): v for k, v in idalabel.items()}
 UpperCAmelCase	:					str	=	idalabel
 UpperCAmelCase	:					Dict	=	{v: k for k, v in idalabel.items()}

 # size of the architecture
 if swiftformer_name == "swiftformer_xs":
  UpperCAmelCase	:					Any	=	[3, 3, 6, 4]
  UpperCAmelCase	:					List[str]	=	[48, 56, 112, 220]

 elif swiftformer_name == "swiftformer_s":
  UpperCAmelCase	:					Dict	=	[3, 3, 9, 6]
  UpperCAmelCase	:					Union[str, Any]	=	[48, 64, 168, 224]

 elif swiftformer_name == "swiftformer_l1":
  UpperCAmelCase	:					int	=	[4, 3, 10, 5]
  UpperCAmelCase	:					Optional[int]	=	[48, 96, 192, 384]

 elif swiftformer_name == "swiftformer_l3":
  UpperCAmelCase	:					Union[str, Any]	=	[4, 4, 12, 6]
  UpperCAmelCase	:					List[Any]	=	[64, 128, 320, 512]

 # load state_dict of original model, remove and rename some keys
 if original_ckpt:
  if original_ckpt.startswith("""https"""						):
   UpperCAmelCase	:					List[str]	=	torch.hub.load_state_dict_from_url(UpperCamelCase			,		map_location="""cpu"""			,		check_hash=UpperCamelCase						)
  else:
   UpperCAmelCase	:					Any	=	torch.load(UpperCamelCase			,		map_location="""cpu"""						)
 UpperCAmelCase	:					Optional[Any]	=	checkpoint

 UpperCAmelCase	:					Dict	=	create_rename_keys(UpperCamelCase						)
 for rename_key_src, rename_key_dest in rename_keys:
  rename_key(UpperCamelCase			,		UpperCamelCase			,		UpperCamelCase						)

 # load HuggingFace model
 UpperCAmelCase	:					List[Any]	=	SwiftFormerForImageClassification(UpperCamelCase						).eval()
 hf_model.load_state_dict(UpperCamelCase						)

 # prepare test inputs
 UpperCAmelCase	:					Dict	=	prepare_img()
 UpperCAmelCase	:					Tuple	=	ViTImageProcessor.from_pretrained("""preprocessor_config"""						)
 UpperCAmelCase	:					Optional[int]	=	processor(images=UpperCamelCase			,		return_tensors="""pt"""						)

 # compare outputs from both models
 UpperCAmelCase	:					Optional[int]	=	get_expected_output(UpperCamelCase						)
 UpperCAmelCase	:					List[str]	=	hf_model(inputs["""pixel_values"""]						).logits

 assert hf_logits.shape == torch.Size([1, 1000]						)
 assert torch.allclose(hf_logits[0, 0:5]			,		UpperCamelCase			,		atol=1e-3						)

 Path(UpperCamelCase						).mkdir(exist_ok=UpperCamelCase						)
 print(F"Saving model {swiftformer_name} to {pytorch_dump_folder_path}"						)
 hf_model.save_pretrained(UpperCamelCase						)


if __name__ == "__main__":
      A:   List[str]								=		argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          "--swiftformer_name",
          default="swiftformer_xs",
          choices=["swiftformer_xs", "swiftformer_s", "swiftformer_l1", "swiftformer_l3"],
          type=str,
          help="Name of the SwiftFormer model you'd like to convert.",
      )
      parser.add_argument(
          "--pytorch_dump_folder_path",
          default="./converted_outputs/",
          type=str,
          help="Path to the output PyTorch model directory.",
      )
      parser.add_argument("--original_ckpt", default=None, type=str, help="Path to the original model checkpoint.")

      A:   str								=		parser.parse_args()
      convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)
