


'''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[int]    						=			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 			__lowerCamelCase			:
    """simple docstring"""


    def __init__(  self      :							List[str]				,							SCREAMING_SNAKE_CASE      :							int				,							SCREAMING_SNAKE_CASE      :							Dict=16				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]=13				,							SCREAMING_SNAKE_CASE      :							List[str]=7				,							SCREAMING_SNAKE_CASE      :							Optional[int]=14				,							SCREAMING_SNAKE_CASE      :							List[str]=10				,							SCREAMING_SNAKE_CASE      :							Dict=19				,							SCREAMING_SNAKE_CASE      :							List[str]=5				,							SCREAMING_SNAKE_CASE      :							str=4				,							SCREAMING_SNAKE_CASE      :							Tuple=True				,							SCREAMING_SNAKE_CASE      :							List[str]=16				,							SCREAMING_SNAKE_CASE      :							Optional[int]=2				,							SCREAMING_SNAKE_CASE      :							int=4				,							SCREAMING_SNAKE_CASE      :							List[str]=4				,							SCREAMING_SNAKE_CASE      :							List[Any]="gelu"				,							SCREAMING_SNAKE_CASE      :							Optional[Any]=0.1				,							SCREAMING_SNAKE_CASE      :							Optional[Any]=0.1				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]=[1, 2, 3, 4, 5]				,							SCREAMING_SNAKE_CASE      :							Tuple=25				,							SCREAMING_SNAKE_CASE      :							Tuple=5				,							):
           _A					:     Tuple     = d_model
           _A					:     Dict     = parent
           _A					:     Optional[Any]     = batch_size
           _A					:     Tuple     = prediction_length
           _A					:     List[str]     = context_length
           _A					:     Dict     = cardinality
           _A					:     List[Any]     = num_time_features
           _A					:     Optional[int]     = lags_sequence
           _A					:     int     = embedding_dimension
           _A					:     Optional[int]     = is_training
           _A					:     Tuple     = hidden_size
           _A					:     Any     = num_hidden_layers
           _A					:     Dict     = num_attention_heads
           _A					:     List[str]     = intermediate_size
           _A					:     Tuple     = hidden_act
           _A					:     List[Any]     = hidden_dropout_prob
           _A					:     List[str]     = attention_probs_dropout_prob

           _A					:     List[Any]     = context_length
           _A					:     int     = prediction_length + label_length
           _A					:     Tuple     = label_length

           _A					:     Optional[int]     = moving_average
           _A					:     Optional[int]     = autocorrelation_factor


    def        A   (  self      :							Dict):
           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        A   (  self      :							Any				,							SCREAMING_SNAKE_CASE      :							Union[str, Any]):
           _A					:     Any     = config.context_length + max(config.lags_sequence)

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

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

           _A					:     int     = {
               '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        A   (  self      :							str):
           _A					:     List[str]     = self.get_config()
           _A					:     Dict     = self.prepare_autoformer_inputs_dict(SCREAMING_SNAKE_CASE)
           return config, inputs_dict


    def        A   (  self      :							List[str]):
           _A  ,			_A					:     Optional[int]     = self.prepare_config_and_inputs()
           return config, inputs_dict


    def        A   (  self      :							int				,							SCREAMING_SNAKE_CASE      :							Optional[int]				,							SCREAMING_SNAKE_CASE      :							List[Any]):
           _A					:     Dict     = AutoformerModel(config=SCREAMING_SNAKE_CASE).to(SCREAMING_SNAKE_CASE).eval()
           _A					:     Any     = model(**SCREAMING_SNAKE_CASE)

           _A					:     Optional[Any]     = outputs.encoder_last_hidden_state
           _A					:     List[Any]     = outputs.last_hidden_state

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

           _A  ,			_A  ,			_A  ,			_A  ,			_A					:     List[Any]     = model.create_network_inputs(**SCREAMING_SNAKE_CASE)
           _A  ,			_A					:     Optional[Any]     = model.decomposition_layer(transformer_inputs[:, : config.context_length, ...])

           _A					:     List[Any]     = torch.cat(
               (transformer_inputs[:, : config.context_length, ...], feature[:, : config.context_length, ...])				,							dim=-1				,							)
           _A					:     List[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)

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

           _A					:     List[Any]     = torch.cat(
               (
                   torch.cat((seasonal_input[:, -config.label_length :, ...], zeros)				,							dim=1),
                   feature[:, config.context_length - config.label_length :, ...],
               )				,							dim=-1				,							)
           _A					:     Tuple     = 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:
                  _A					:     Optional[int]     = model.get_decoder()
                  decoder.save_pretrained(SCREAMING_SNAKE_CASE)
                  _A					:     Any     = AutoformerDecoder.from_pretrained(SCREAMING_SNAKE_CASE).to(SCREAMING_SNAKE_CASE)

           _A					:     int     = 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 			__lowerCamelCase			(	a_    , a_    , unittest.TestCase	):
    """simple docstring"""
    a													=   (AutoformerModel, AutoformerForPrediction) if is_torch_available() else ()
    a													=   (AutoformerForPrediction,) if is_torch_available() else ()
    a													=   {"feature-extraction": AutoformerModel} if is_torch_available() else {}
    a													=   False
    a													=   False
    a													=   False
    a													=   False
    a													=   False
    a													=   False


    def        A   (  self      :							List[str]):
           _A					:     Optional[int]     = AutoformerModelTester(self)
           _A					:     Union[str, Any]     = ConfigTester(self				,							config_class=SCREAMING_SNAKE_CASE				,							has_text_modality=SCREAMING_SNAKE_CASE)


    def        A   (  self      :							List[Any]):
           self.config_tester.run_common_tests()


    def        A   (  self      :							Union[str, Any]):
           _A  ,			_A					:     List[str]     = self.model_tester.prepare_config_and_inputs()
           for model_class in self.all_model_classes:
                  _A					:     Any     = model_class(SCREAMING_SNAKE_CASE)

                  with tempfile.TemporaryDirectory() as tmpdirname:
                         model.save_pretrained(SCREAMING_SNAKE_CASE)
                         _A  ,			_A					:     str     = model_class.from_pretrained(SCREAMING_SNAKE_CASE				,							output_loading_info=SCREAMING_SNAKE_CASE)
                  self.assertEqual(info['missing_keys']				,							[])


    def        A   (  self      :							int):
           _A					:     Dict     = 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        A   (  self      :							Any):
           pass


    def        A   (  self      :							List[str]):
           _A					:     Dict     = inspect.signature(getattr(SCREAMING_SNAKE_CASE				,							'forward'))
           # The main input is the name of the argument after `self`
           _A					:     str     = list(model_signature.parameters.keys())[1]
           self.assertEqual(AutoformerModel.main_input_name				,							SCREAMING_SNAKE_CASE)


    def        A   (  self      :							Optional[int]):
           _A  ,			_A					:     int     = self.model_tester.prepare_config_and_inputs_for_common()

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

                  _A					:     str     = [
                      '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        A   (  self      :							Union[str, Any]):
           _A  ,			_A					:     List[str]     = self.model_tester.prepare_config_and_inputs_for_common()
           _A					:     int     = True

           _A					:     Optional[int]     = getattr(self.model_tester				,							'seq_length'				,							SCREAMING_SNAKE_CASE)
           _A					:     Dict     = getattr(self.model_tester				,							'decoder_seq_length'				,							SCREAMING_SNAKE_CASE)
           _A					:     Any     = getattr(self.model_tester				,							'encoder_seq_length'				,							SCREAMING_SNAKE_CASE)
           _A					:     Optional[Any]     = getattr(self.model_tester				,							'd_model'				,							SCREAMING_SNAKE_CASE)
           _A					:     Tuple     = getattr(self.model_tester				,							'num_attention_heads'				,							SCREAMING_SNAKE_CASE)
           _A					:     int     = d_model // num_attention_heads

           for model_class in self.all_model_classes:
                  _A					:     str     = True
                  _A					:     int     = False
                  _A					:     Tuple     = True
                  _A					:     str     = model_class(SCREAMING_SNAKE_CASE)
                  model.to(SCREAMING_SNAKE_CASE)
                  model.eval()
                  with torch.no_grad():
                         _A					:     Tuple     = model(**self._prepare_for_class(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE))
                  _A					:     Any     = 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"]
                  _A					:     List[str]     = True
                  _A					:     List[Any]     = model_class(SCREAMING_SNAKE_CASE)
                  model.to(SCREAMING_SNAKE_CASE)
                  model.eval()
                  with torch.no_grad():
                         _A					:     Any     = model(**self._prepare_for_class(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE))
                  _A					:     Optional[int]     = 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]				,							)
                  _A					:     Union[str, Any]     = len(SCREAMING_SNAKE_CASE)

                  _A					:     Tuple     = 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
                  _A					:     List[Any]     = 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
                  _A					:     Optional[int]     = 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
           _A					:     Tuple     = True
           _A					:     List[Any]     = True
           _A					:     int     = model_class(SCREAMING_SNAKE_CASE)
           model.to(SCREAMING_SNAKE_CASE)
           model.eval()
           with torch.no_grad():
                  _A					:     Tuple     = model(**self._prepare_for_class(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE))

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

           _A					:     int     = 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        A   (  self      :							int):
           super().test_retain_grad_hidden_states_attentions()






def 				lowerCAmelCase__	(    lowerCamelCase					:      Union[str, Any]="train-batch.pt"       ):
       _A					:     Any     = hf_hub_download(repo_id='hf-internal-testing/tourism-monthly-batch'  ,filename=lowerCamelCase  ,repo_type='dataset'       )
       _A					:     Optional[Any]     = torch.load(lowerCamelCase  ,map_location=lowerCamelCase       )
       return batch


@require_torch
@slow
class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""


    def        A   (  self      :							int):
           _A					:     Any     = AutoformerModel.from_pretrained('huggingface/autoformer-tourism-monthly').to(SCREAMING_SNAKE_CASE)
           _A					:     int     = prepare_batch()

           with torch.no_grad():
                  _A					:     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']				,							future_values=batch['future_values']				,							future_time_features=batch['future_time_features']				,							)[0]

           _A					:     Any     = torch.Size(
               (64, model.config.prediction_length + model.config.label_length, model.config.feature_size))
           self.assertEqual(output.shape				,							SCREAMING_SNAKE_CASE)

           _A					:     Union[str, Any]     = 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        A   (  self      :							Dict):
           _A					:     Optional[int]     = AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly').to(SCREAMING_SNAKE_CASE)
           _A					:     Any     = prepare_batch('val-batch.pt')
           with torch.no_grad():
                  _A					:     Tuple     = 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
           _A					:     Optional[int]     = torch.Size((64, model.config.context_length, model.config.d_model))
           self.assertEqual(output.shape				,							SCREAMING_SNAKE_CASE)

           _A					:     List[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        A   (  self      :							str):
           _A					:     int     = AutoformerForPrediction.from_pretrained('huggingface/autoformer-tourism-monthly').to(SCREAMING_SNAKE_CASE)
           _A					:     int     = prepare_batch('val-batch.pt')
           with torch.no_grad():
                  _A					:     str     = 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']				,							)
           _A					:     int     = torch.Size((64, model.config.num_parallel_samples, model.config.prediction_length))
           self.assertEqual(outputs.sequences.shape				,							SCREAMING_SNAKE_CASE)

           _A					:     Union[str, Any]     = torch.tensor([3130.6763, 4056.5293, 7053.0786]				,							device=SCREAMING_SNAKE_CASE)
           _A					:     List[Any]     = outputs.sequences.mean(dim=1)
           self.assertTrue(torch.allclose(mean_prediction[0, -3:]				,							SCREAMING_SNAKE_CASE				,							rtol=1e-1))










'''simple docstring'''







import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)






def 				lowerCAmelCase__	(    lowerCamelCase					:      Optional[int]       ):
       if model.config.model_type == "gpt2":
              return model.transformer.h[0].mlp.c_fc
       return model.transformer.h[0].mlp.dense_ah_to_h


if is_torch_available():
   import torch
   import torch.nn as nn


   class 			__lowerCamelCase			(	nn.Module	):
       """simple docstring"""


       def __init__(  self      :							int				,							SCREAMING_SNAKE_CASE      :							nn.Module				,							SCREAMING_SNAKE_CASE      :							int):
              super().__init__()
              _A					:     Tuple     = module
              _A					:     str     = nn.Sequential(
                  nn.Linear(module.in_features				,							SCREAMING_SNAKE_CASE				,							bias=SCREAMING_SNAKE_CASE)				,							nn.Linear(SCREAMING_SNAKE_CASE				,							module.out_features				,							bias=SCREAMING_SNAKE_CASE)				,							)
              _A					:     Any     = (2.0 / (5 * min(module.in_features				,							module.out_features))) ** 0.5
              nn.init.normal_(self.adapter[0].weight				,							std=SCREAMING_SNAKE_CASE)
              nn.init.zeros_(self.adapter[1].weight)
              self.adapter.to(module.weight.device)


       def        A   (  self      :							Any				,							SCREAMING_SNAKE_CASE      :							Any				,							*SCREAMING_SNAKE_CASE      :							List[str]				,							**SCREAMING_SNAKE_CASE      :							int):
              return self.module(SCREAMING_SNAKE_CASE				,							*SCREAMING_SNAKE_CASE				,							**SCREAMING_SNAKE_CASE) + self.adapter(SCREAMING_SNAKE_CASE)


@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""
    # We keep the constants inside the init function and model loading inside setUp function

    # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
    # Therefore here we use only bloom-1b3 to test our module
    a													=   "bigscience/bloom-1b7"

    # Constant values
    a													=   2.109_6595_5269_2574

    a													=   "Hello my name is"
    a													=   set()
    EXPECTED_OUTPUTS.add("Hello my name is John and I am a professional photographer. I"	)
    EXPECTED_OUTPUTS.add("Hello my name is John.\nI am a friend of your father.\n"	)
    EXPECTED_OUTPUTS.add("Hello my name is John Doe, I am a student at the University"	)
    a													=   10


    def        A   (  self      :							List[Any]):
           # Models and tokenizer
           _A					:     List[Any]     = AutoTokenizer.from_pretrained(self.model_name)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							List[str]):
           super().setUp()

           # Models and tokenizer
           _A					:     Tuple     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							torch_dtype=torch.floataa				,							device_map='auto')
           _A					:     str     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')


    def        A   (  self      :							List[Any]):
           del self.model_fpaa
           del self.model_abit

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							str):
           _A					:     Any     = self.model_abit.config

           self.assertTrue(hasattr(SCREAMING_SNAKE_CASE				,							'quantization_config'))

           _A					:     str     = config.to_dict()
           _A					:     Dict     = config.to_diff_dict()

           _A					:     int     = config.to_json_string()


    def        A   (  self      :							List[str]):
           from bitsandbytes.nn import Paramsabit

           _A					:     List[str]     = self.model_fpaa.get_memory_footprint()
           _A					:     List[Any]     = self.model_abit.get_memory_footprint()

           self.assertAlmostEqual(mem_fpaa / mem_abit				,							self.EXPECTED_RELATIVE_DIFFERENCE)
           _A					:     str     = get_some_linear_layer(self.model_abit)
           self.assertTrue(linear.weight.__class__ == Paramsabit)


    def        A   (  self      :							List[Any]):
           from transformers import TaPreTrainedModel

           self.model_fpaa.get_memory_footprint()
           self.model_abit.get_memory_footprint()

           for name, module in self.model_abit.named_modules():
                  if isinstance(SCREAMING_SNAKE_CASE				,							torch.nn.Linear):
                         if name not in ["lm_head"] + TaPreTrainedModel._keep_in_fpaa_modules:
                                # 4-bit parameters are packed in uint8 variables
                                self.assertTrue(module.weight.dtype == torch.uinta)


    def        A   (  self      :							Any):
           _A					:     List[str]     = self.tokenizer(self.input_text				,							return_tensors='pt')
           _A					:     Any     = self.model_abit.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           self.assertIn(self.tokenizer.decode(output_sequences[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)


    def        A   (  self      :							List[str]):
           _A					:     List[Any]     = BitsAndBytesConfig()
           _A					:     Optional[int]     = True

           _A					:     Union[str, Any]     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							quantization_config=SCREAMING_SNAKE_CASE				,							device_map='auto')

           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt')
           _A					:     Any     = model_abit_from_config.generate(
               input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           self.assertIn(self.tokenizer.decode(output_sequences[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)


    def        A   (  self      :							Union[str, Any]):
           with self.assertRaises(SCREAMING_SNAKE_CASE), tempfile.TemporaryDirectory() as tmpdirname:
                  self.model_abit.save_pretrained(SCREAMING_SNAKE_CASE)


    def        A   (  self      :							int):
           _A					:     Any     = BitsAndBytesConfig()

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  _A					:     Optional[Any]     = AutoModelForCausalLM.from_pretrained(
                      self.model_name				,							quantization_config=SCREAMING_SNAKE_CASE				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto'				,							bnb_abit_quant_type='nf4'				,							)


    def        A   (  self      :							Optional[Any]):
           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with `str`
                  self.model_abit.to('cpu')

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `dtype``
                  self.model_abit.to(torch.floataa)

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.to(torch.device('cuda:0'))

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.float()

           with self.assertRaises(SCREAMING_SNAKE_CASE):
                  # Tries with a `device`
                  self.model_abit.half()

           # Test if we did not break anything
           _A					:     str     = self.tokenizer(self.input_text				,							return_tensors='pt')

           _A					:     Optional[Any]     = self.model_fpaa.to(torch.floataa)
           _A					:     Optional[int]     = self.model_fpaa.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)

           # Check this does not throw an error
           _A					:     Tuple     = self.model_fpaa.to('cpu')

           # Check this does not throw an error
           _A					:     Optional[int]     = self.model_fpaa.half()

           # Check this does not throw an error
           _A					:     int     = self.model_fpaa.float()


    def        A   (  self      :							Dict):
           _A					:     Any     = AutoModelForSeqaSeqLM.from_pretrained('t5-small'				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           self.assertTrue(model.decoder.block[0].layer[2].DenseReluDense.wo.weight.dtype == torch.floataa)




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__lowerCamelCase			(	unittest.TestCase	):
    """simple docstring"""


    @classmethod
    def        A   (  cls      :							Union[str, Any]):
           _A					:     Tuple     = 't5-small'
           _A					:     Union[str, Any]     = 'google/flan-t5-small'  # flan-t5 uses dense-act instead of dense-relu-dense
           _A					:     str     = AutoTokenizer.from_pretrained(cls.model_name)
           _A					:     Optional[int]     = 'Translate in German: Hello, my dog is cute'


    def        A   (  self      :							Any):
           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							List[str]):
           from transformers import TaForConditionalGeneration

           _A					:     Union[str, Any]     = TaForConditionalGeneration._keep_in_fpaa_modules
           _A					:     List[str]     = None

           # test with `t5-small`
           _A					:     Optional[int]     = TaForConditionalGeneration.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     List[str]     = model.generate(**SCREAMING_SNAKE_CASE)

           # test with `flan-t5-small`
           _A					:     List[str]     = TaForConditionalGeneration.from_pretrained(
               self.dense_act_model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     Dict     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Any     = model.generate(**SCREAMING_SNAKE_CASE)
           _A					:     List[str]     = modules


    def        A   (  self      :							Dict):
           import bitsandbytes as bnb

           from transformers import TaForConditionalGeneration

           # test with `t5-small`
           _A					:     str     = TaForConditionalGeneration.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')

           # there was a bug with decoders - this test checks that it is fixed
           self.assertTrue(isinstance(model.decoder.block[0].layer[0].SelfAttention.q				,							bnb.nn.Linearabit))

           _A					:     int     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Union[str, Any]     = model.generate(**SCREAMING_SNAKE_CASE)

           # test with `flan-t5-small`
           _A					:     Union[str, Any]     = TaForConditionalGeneration.from_pretrained(
               self.dense_act_model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           _A					:     str     = self.tokenizer(self.input_text				,							return_tensors='pt').to(0)
           _A					:     Optional[int]     = model.generate(**SCREAMING_SNAKE_CASE)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Union[str, Any]):
           super().setUp()
           # model_name
           _A					:     Dict     = 'bigscience/bloom-560m'
           _A					:     List[str]     = 't5-small'

           # Different types of model

           _A					:     Optional[Any]     = AutoModel.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # Sequence classification model
           _A					:     Tuple     = AutoModelForSequenceClassification.from_pretrained(
               self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # CausalLM model
           _A					:     List[Any]     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')
           # Seq2seq model
           _A					:     Dict     = AutoModelForSeqaSeqLM.from_pretrained(
               self.seq_to_seq_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='auto')


    def        A   (  self      :							str):
           del self.base_model
           del self.sequence_model
           del self.model_abit
           del self.seq_to_seq_model

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							Any):
           from bitsandbytes.nn import Paramsabit

           self.assertTrue(self.base_model.h[-1].mlp.dense_ah_to_h.weight.__class__ == Paramsabit)

           # Other heads should be nn.Parameter
           self.assertTrue(self.model_abit.lm_head.weight.__class__ == torch.nn.Parameter)
           self.assertTrue(self.sequence_model.score.weight.__class__ == torch.nn.Parameter)
           self.assertTrue(self.seq_to_seq_model.lm_head.weight.__class__ == torch.nn.Parameter)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Optional[Any]):
           super().setUp()


    def        A   (  self      :							Any):
           del self.pipe

           gc.collect()
           torch.cuda.empty_cache()


    def        A   (  self      :							Optional[int]):
           _A					:     Dict     = pipeline(
               'text-generation'				,							model=self.model_name				,							model_kwargs={'device_map': 'auto', 'load_in_4bit': True, 'torch_dtype': torch.floataa}				,							max_new_tokens=self.MAX_NEW_TOKENS				,							)

           # Real second forward pass
           _A					:     Union[str, Any]     = self.pipe(self.input_text)
           self.assertIn(pipeline_output[0]['generated_text']				,							self.EXPECTED_OUTPUTS)




@require_torch_multi_gpu
class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							List[Any]):
           super().setUp()


    def        A   (  self      :							Any):
           _A					:     int     = AutoModelForCausalLM.from_pretrained(
               self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE				,							device_map='balanced')

           # Check correct device map
           self.assertEqual(set(model_parallel.hf_device_map.values())				,							{0, 1})

           # Check that inference pass works on the model
           _A					:     Tuple     = self.tokenizer(self.input_text				,							return_tensors='pt')

           # Second real batch
           _A					:     Union[str, Any]     = model_parallel.generate(input_ids=encoded_input['input_ids'].to(0)				,							max_new_tokens=10)
           self.assertIn(self.tokenizer.decode(output_parallel[0]				,							skip_special_tokens=SCREAMING_SNAKE_CASE)				,							self.EXPECTED_OUTPUTS)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""


    def        A   (  self      :							Optional[Any]):
           _A					:     Any     = 'facebook/opt-350m'
           super().setUp()


    def        A   (  self      :							List[Any]):
           if version.parse(importlib.metadata.version('bitsandbytes')) < version.parse('0.37.0'):
                  return

           # Step 1: freeze all parameters
           _A					:     Any     = AutoModelForCausalLM.from_pretrained(self.model_name				,							load_in_abit=SCREAMING_SNAKE_CASE)

           self.assertEqual(set(model.hf_device_map.values())				,							{torch.cuda.current_device()})

           for param in model.parameters():
                  _A					:     Optional[Any]     = False  # freeze the model - train adapters later
                  if param.ndim == 1:
                         # cast the small parameters (e.g. layernorm) to fp32 for stability
                         _A					:     Tuple     = param.data.to(torch.floataa)

        # Step 2: add adapters
           for _, module in model.named_modules():
                  if "OPTAttention" in repr(type(SCREAMING_SNAKE_CASE)):
                         _A					:     List[Any]     = LoRALayer(module.q_proj				,							rank=16)
                         _A					:     Union[str, Any]     = LoRALayer(module.k_proj				,							rank=16)
                         _A					:     str     = LoRALayer(module.v_proj				,							rank=16)

        # Step 3: dummy batch
           _A					:     Union[str, Any]     = self.tokenizer('Test batch '				,							return_tensors='pt').to(0)

           # Step 4: Check if the gradient is not None
           with torch.cuda.amp.autocast():
                  _A					:     Union[str, Any]     = model.forward(**SCREAMING_SNAKE_CASE)
                  out.logits.norm().backward()

           for module in model.modules():
                  if isinstance(SCREAMING_SNAKE_CASE				,							SCREAMING_SNAKE_CASE):
                         self.assertTrue(module.adapter[1].weight.grad is not None)
                         self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0)
                  elif isinstance(SCREAMING_SNAKE_CASE				,							nn.Embedding):
                         self.assertTrue(module.weight.grad is None)




class 			__lowerCamelCase			(	a_	):
    """simple docstring"""
    a													=   "gpt2-xl"
    a													=   3.3191_8548_5415_2187






