



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   (						A__	):
      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       UpperCAmelCase__(							nn.Module	):

          '''simple docstring'''




          def __init__(						self  :			Any			,				lowerCAmelCase  :			nn.Module			,				lowerCAmelCase  :			int)	->				Dict:




                """simple docstring"""

                super().__init__()
                lowercase__     =  module
                lowercase__     =  nn.Sequential(
                    nn.Linear(module.in_features			,				lowerCAmelCase			,				bias=lowerCAmelCase)			,				nn.Linear(lowerCAmelCase			,				module.out_features			,				bias=lowerCAmelCase)			,				)
                lowercase__     =  (2.0 / (5 * min(module.in_features			,				module.out_features))) ** 0.5
                nn.init.normal_(self.adapter[0].weight			,				std=lowerCAmelCase)
                nn.init.zeros_(self.adapter[1].weight)
                self.adapter.to(module.weight.device)



          def 			UpperCAmelCase		(						self  :			str			,				lowerCAmelCase  :			Dict			,				*lowerCAmelCase  :			Tuple			,				**lowerCAmelCase  :			Tuple)	->				List[str]:




                """simple docstring"""

                return self.module(lowerCAmelCase			,				*lowerCAmelCase			,				**lowerCAmelCase) + self.adapter(lowerCAmelCase)



@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class       UpperCAmelCase__(							unittest.TestCase	):

   '''simple docstring'''




   A    :  Optional[int]       				=    "bigscience/bloom-1b7"

   # Constant values
   A    :  List[str]       				=    2.109_6595_5269_2574

   A    :  Any       				=    "Hello my name is"
   A    :  List[str]       				=    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    :  List[Any]       				=    10


   def 			UpperCAmelCase		(						self  :			int)	->				Optional[Any]:




         """simple docstring"""

         lowercase__     =  AutoTokenizer.from_pretrained(self.model_name)



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




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




         """simple docstring"""

         super().setUp()

         # Models and tokenizer
         lowercase__     =  AutoModelForCausalLM.from_pretrained(
             self.model_name			,				torch_dtype=torch.floataa			,				device_map='auto')
         lowercase__     =  AutoModelForCausalLM.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')


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




         """simple docstring"""

         del self.model_fpaa
         del self.model_abit

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


   def 			UpperCAmelCase		(						self  :			Tuple)	->				str:




         """simple docstring"""

         lowercase__     =  self.model_abit.config

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

         lowercase__     =  config.to_dict()
         lowercase__     =  config.to_diff_dict()

         lowercase__     =  config.to_json_string()


   def 			UpperCAmelCase		(						self  :			List[str])	->				Any:




         """simple docstring"""

         from bitsandbytes.nn import Paramsabit

         lowercase__     =  self.model_fpaa.get_memory_footprint()
         lowercase__     =  self.model_abit.get_memory_footprint()

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


   def 			UpperCAmelCase		(						self  :			str)	->				Tuple:




         """simple docstring"""

         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(lowerCAmelCase			,				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 			UpperCAmelCase		(						self  :			List[Any])	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt')
         lowercase__     =  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=lowerCAmelCase)			,				self.EXPECTED_OUTPUTS)


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




         """simple docstring"""

         lowercase__     =  BitsAndBytesConfig()
         lowercase__     =  True

         lowercase__     =  AutoModelForCausalLM.from_pretrained(
             self.model_name			,				quantization_config=lowerCAmelCase			,				device_map='auto')

         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt')
         lowercase__     =  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=lowerCAmelCase)			,				self.EXPECTED_OUTPUTS)


   def 			UpperCAmelCase		(						self  :			Dict)	->				List[str]:




         """simple docstring"""

         with self.assertRaises(lowerCAmelCase), tempfile.TemporaryDirectory() as tmpdirname:
               self.model_abit.save_pretrained(lowerCAmelCase)


   def 			UpperCAmelCase		(						self  :			List[Any])	->				Dict:




         """simple docstring"""

         lowercase__     =  BitsAndBytesConfig()

         with self.assertRaises(lowerCAmelCase):
               lowercase__     =  AutoModelForCausalLM.from_pretrained(
                   self.model_name			,				quantization_config=lowerCAmelCase			,				load_in_abit=lowerCAmelCase			,				device_map='auto'			,				bnb_abit_quant_type='nf4'			,				)


   def 			UpperCAmelCase		(						self  :			Dict)	->				List[str]:




         """simple docstring"""

         with self.assertRaises(lowerCAmelCase):
               # Tries with `str`
               self.model_abit.to('cpu')

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

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

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

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

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

         lowercase__     =  self.model_fpaa.to(torch.floataa)
         lowercase__     =  self.model_fpaa.generate(input_ids=encoded_input['input_ids'].to(0)			,				max_new_tokens=10)

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

         # Check this does not throw an error
         lowercase__     =  self.model_fpaa.half()

         # Check this does not throw an error
         lowercase__     =  self.model_fpaa.float()



   def 			UpperCAmelCase		(						self  :			str)	->				Tuple:




         """simple docstring"""

         lowercase__     =  AutoModelForSeqaSeqLM.from_pretrained('t5-small'			,				load_in_abit=lowerCAmelCase			,				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       UpperCAmelCase__(							unittest.TestCase	):

   '''simple docstring'''




   @classmethod
   def 			UpperCAmelCase		(						cls  :			Optional[Any])	->				Tuple:




         """simple docstring"""

         lowercase__     =  't5-small'
         lowercase__     =  'google/flan-t5-small'  # flan-t5 uses dense-act instead of dense-relu-dense
         lowercase__     =  AutoTokenizer.from_pretrained(cls.model_name)
         lowercase__     =  'Translate in German: Hello, my dog is cute'


   def 			UpperCAmelCase		(						self  :			str)	->				Optional[int]:




         """simple docstring"""

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


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




         """simple docstring"""

         from transformers import TaForConditionalGeneration

         lowercase__     =  TaForConditionalGeneration._keep_in_fpaa_modules
         lowercase__     =  None

         # test with `t5-small`
         lowercase__     =  TaForConditionalGeneration.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt').to(0)
         lowercase__     =  model.generate(**lowerCAmelCase)

         # test with `flan-t5-small`
         lowercase__     =  TaForConditionalGeneration.from_pretrained(
             self.dense_act_model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt').to(0)
         lowercase__     =  model.generate(**lowerCAmelCase)
         lowercase__     =  modules



   def 			UpperCAmelCase		(						self  :			List[Any])	->				int:




         """simple docstring"""

         import bitsandbytes as bnb

         from transformers import TaForConditionalGeneration

         # test with `t5-small`
         lowercase__     =  TaForConditionalGeneration.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase			,				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))

         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt').to(0)
         lowercase__     =  model.generate(**lowerCAmelCase)

         # test with `flan-t5-small`
         lowercase__     =  TaForConditionalGeneration.from_pretrained(
             self.dense_act_model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt').to(0)
         lowercase__     =  model.generate(**lowerCAmelCase)



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def 			UpperCAmelCase		(						self  :			List[Any])	->				int:




         """simple docstring"""

         super().setUp()
         # model_name
         lowercase__     =  'bigscience/bloom-560m'
         lowercase__     =  't5-small'

         # Different types of model

         lowercase__     =  AutoModel.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         # Sequence classification model
         lowercase__     =  AutoModelForSequenceClassification.from_pretrained(
             self.model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         # CausalLM model
         lowercase__     =  AutoModelForCausalLM.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')
         # Seq2seq model
         lowercase__     =  AutoModelForSeqaSeqLM.from_pretrained(
             self.seq_to_seq_name			,				load_in_abit=lowerCAmelCase			,				device_map='auto')


   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				Dict:




         """simple docstring"""

         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 			UpperCAmelCase		(						self  :			Optional[Any])	->				str:




         """simple docstring"""

         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       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def 			UpperCAmelCase		(						self  :			Optional[int])	->				Optional[Any]:




         """simple docstring"""

         super().setUp()


   def 			UpperCAmelCase		(						self  :			Any)	->				Optional[Any]:




         """simple docstring"""

         del self.pipe

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



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




         """simple docstring"""

         lowercase__     =  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
         lowercase__     =  self.pipe(self.input_text)
         self.assertIn(pipeline_output[0]['generated_text']			,				self.EXPECTED_OUTPUTS)



@require_torch_multi_gpu
class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def 			UpperCAmelCase		(						self  :			Optional[Any])	->				List[str]:




         """simple docstring"""

         super().setUp()



   def 			UpperCAmelCase		(						self  :			List[Any])	->				Optional[int]:




         """simple docstring"""

         lowercase__     =  AutoModelForCausalLM.from_pretrained(
             self.model_name			,				load_in_abit=lowerCAmelCase			,				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
         lowercase__     =  self.tokenizer(self.input_text			,				return_tensors='pt')

         # Second real batch
         lowercase__     =  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=lowerCAmelCase)			,				self.EXPECTED_OUTPUTS)



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   def 			UpperCAmelCase		(						self  :			List[Any])	->				Dict:




         """simple docstring"""

         lowercase__     =  'facebook/opt-350m'
         super().setUp()



   def 			UpperCAmelCase		(						self  :			str)	->				Tuple:




         """simple docstring"""

         if version.parse(importlib.metadata.version('bitsandbytes')) < version.parse('0.37.0'):
               return

         # Step 1: freeze all parameters
         lowercase__     =  AutoModelForCausalLM.from_pretrained(self.model_name			,				load_in_abit=lowerCAmelCase)

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

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

        # Step 2: add adapters
         for _, module in model.named_modules():
               if "OPTAttention" in repr(type(lowerCAmelCase)):
                     lowercase__     =  LoRALayer(module.q_proj			,				rank=16)
                     lowercase__     =  LoRALayer(module.k_proj			,				rank=16)
                     lowercase__     =  LoRALayer(module.v_proj			,				rank=16)

        # Step 3: dummy batch
         lowercase__     =  self.tokenizer('Test batch '			,				return_tensors='pt').to(0)

         # Step 4: Check if the gradient is not None
         with torch.cuda.amp.autocast():
               lowercase__     =  model.forward(**lowerCAmelCase)
               out.logits.norm().backward()

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



class       UpperCAmelCase__(							lowerCamelCase	):

   '''simple docstring'''




   A    :  List[str]       				=    "gpt2-xl"
   A    :  Any       				=    3.3191_8548_5415_2187





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a__      :				List[str]   	=			logging.get_logger(__name__)

a__      :				List[Any]   	=			{
    "microsoft/focalnet-tiny": "https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json",
}



class       UpperCAmelCase__(							lowerCamelCase     ,							lowerCamelCase	):

   '''simple docstring'''




   A    :  List[str]       				=    "focalnet"


   def __init__(						self  :			Dict			,				lowerCAmelCase  :			Union[str, Any]=2_24			,				lowerCAmelCase  :			List[str]=4			,				lowerCAmelCase  :			int=3			,				lowerCAmelCase  :			Union[str, Any]=96			,				lowerCAmelCase  :			List[Any]=False			,				lowerCAmelCase  :			int=[1_92, 3_84, 7_68, 7_68]			,				lowerCAmelCase  :			str=[2, 2, 6, 2]			,				lowerCAmelCase  :			Tuple=[2, 2, 2, 2]			,				lowerCAmelCase  :			Optional[Any]=[3, 3, 3, 3]			,				lowerCAmelCase  :			int="gelu"			,				lowerCAmelCase  :			Any=4.0			,				lowerCAmelCase  :			List[str]=0.0			,				lowerCAmelCase  :			Union[str, Any]=0.1			,				lowerCAmelCase  :			Union[str, Any]=False			,				lowerCAmelCase  :			Tuple=1E-4			,				lowerCAmelCase  :			List[Any]=False			,				lowerCAmelCase  :			Optional[int]=False			,				lowerCAmelCase  :			List[str]=False			,				lowerCAmelCase  :			str=0.02			,				lowerCAmelCase  :			Optional[int]=1E-5			,				lowerCAmelCase  :			List[Any]=32			,				lowerCAmelCase  :			List[Any]=None			,				lowerCAmelCase  :			Union[str, Any]=None			,				**lowerCAmelCase  :			str			,				)	->				List[str]:




         """simple docstring"""

         super().__init__(**lowerCAmelCase)

         lowercase__     =  image_size
         lowercase__     =  patch_size
         lowercase__     =  num_channels
         lowercase__     =  embed_dim
         lowercase__     =  use_conv_embed
         lowercase__     =  hidden_sizes
         lowercase__     =  depths
         lowercase__     =  focal_levels
         lowercase__     =  focal_windows
         lowercase__     =  hidden_act
         lowercase__     =  mlp_ratio
         lowercase__     =  hidden_dropout_prob
         lowercase__     =  drop_path_rate
         lowercase__     =  use_layerscale
         lowercase__     =  layerscale_value
         lowercase__     =  use_post_layernorm
         lowercase__     =  use_post_layernorm_in_modulation
         lowercase__     =  normalize_modulator
         lowercase__     =  initializer_range
         lowercase__     =  layer_norm_eps
         lowercase__     =  encoder_stride
         lowercase__     =  ['stem'] + [f'''stage{idx}''' for idx in range(1			,				len(self.depths) + 1)]
         lowercase__,	lowercase__     =  get_aligned_output_features_output_indices(
             out_features=lowerCAmelCase			,				out_indices=lowerCAmelCase			,				stage_names=self.stage_names)
