



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      __snake_case	(			_UpperCAmelCase				):
    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     _A       (	nn.Module    ):



        def __init__(		self							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     nn.Module     ,		__SCREAMING_SNAKE_CASE							:     int):


            '''simple docstring'''

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





        def 		_lowerCamelCase			(		self							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     List[Any]     ,		*__SCREAMING_SNAKE_CASE							:     int     ,		**__SCREAMING_SNAKE_CASE							:     List[Any]):


            '''simple docstring'''

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

@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class     _A       (	unittest.TestCase    ):
       # 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
       UpperCamelCase__  :    Union[str, Any]											=						'''bigscience/bloom-1b7'''

       # Constant values
       UpperCamelCase__  :    Optional[int]											=						2.109659552692574

       UpperCamelCase__  :    List[str]											=						'''Hello my name is'''
       UpperCamelCase__  :    Tuple											=						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'''    )
       UpperCamelCase__  :    Tuple											=						10



       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           __a						       =				AutoTokenizer.from_pretrained(self.model_name)


class     _A       (	lowercase_    ):



       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

           super().setUp()

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



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           del self.model_fpaa
           del self.model_abit

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



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           __a						       =				self.model_abit.config

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

           __a						       =				config.to_dict()
           __a						       =				config.to_diff_dict()

           __a						       =				config.to_json_string()



       def 		_lowerCamelCase			(		self							:     int):


           '''simple docstring'''

           from bitsandbytes.nn import Paramsabit

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

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



       def 		_lowerCamelCase			(		self							:     int):


           '''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(__UpperCamelCase     ,		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 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           __a						       =				BitsAndBytesConfig()
           __a						       =				True

           __a						       =				AutoModelForCausalLM.from_pretrained(
               self.model_name     ,		quantization_config=__UpperCamelCase     ,		device_map='''auto''')

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



       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Optional[Any]):


           '''simple docstring'''

           __a						       =				BitsAndBytesConfig()

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



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

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

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

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

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

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

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

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

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

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

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





       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

           __a						       =				AutoModelForSeqaSeqLM.from_pretrained('''t5-small'''     ,		load_in_abit=__UpperCamelCase     ,		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     _A       (	unittest.TestCase    ):



       @classmethod
       def 		_lowerCamelCase			(		cls							:     Optional[Any]):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

           from transformers import TaForConditionalGeneration

           __a						       =				TaForConditionalGeneration._keep_in_fpaa_modules
           __a						       =				None

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

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





       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           import bitsandbytes as bnb

           from transformers import TaForConditionalGeneration

           # test with `t5-small`
           __a						       =				TaForConditionalGeneration.from_pretrained(self.model_name     ,		load_in_abit=__UpperCamelCase     ,		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						       =				self.tokenizer(self.input_text     ,		return_tensors='''pt''').to(0)
           __a						       =				model.generate(**__UpperCamelCase)

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


class     _A       (	lowercase_    ):



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

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

           # Different types of model

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



       def 		_lowerCamelCase			(		self							:     Any):


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


           '''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     _A       (	lowercase_    ):



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

           super().setUp()



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           del self.pipe

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





       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

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


@require_torch_multi_gpu
class     _A       (	lowercase_    ):



       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

           super().setUp()





       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

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

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


class     _A       (	lowercase_    ):



       def 		_lowerCamelCase			(		self							:     Optional[Any]):


           '''simple docstring'''

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





       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

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

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

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

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

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

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

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

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


class     _A       (	lowercase_    ):
       UpperCamelCase__  :    str											=						'''gpt2-xl'''
       UpperCamelCase__  :    List[Any]											=						3.3191854854152187












"""simple docstring"""
from typing import Callable, Dict, Optional, Tuple

import torch
from torch import nn
from torch.distributions import (
    AffineTransform,
    Distribution,
    Independent,
    NegativeBinomial,
    Normal,
    StudentT,
    TransformedDistribution,
)






class 					_UpperCAmelCase   (	lowercase_ ):


					def __init__(					self			:int     ,       __UpperCamelCase			:Distribution     ,       __UpperCamelCase			:Dict=None     ,       __UpperCamelCase			:Optional[int]=None     ,       __UpperCamelCase			:List[str]=0		):
												A		=  1.0 if scale is None else scale
												A		=  0.0 if loc is None else loc

												super().__init__(__UpperCamelCase     ,       [AffineTransform(loc=self.loc     ,       scale=self.scale     ,       event_dim=__UpperCamelCase		)]		)


					@property
					def 	lowerCamelCase       (					self			:Any		):
												return self.base_dist.mean * self.scale + self.loc


					@property
					def 	lowerCamelCase       (					self			:Optional[int]		):
												return self.base_dist.variance * self.scale**2


					@property
					def 	lowerCamelCase       (					self			:Dict		):
												return self.variance.sqrt()







class 					_UpperCAmelCase   (	nn.Module ):


					def __init__(					self			:Dict     ,       __UpperCamelCase			:int     ,       __UpperCamelCase			:Dict[str, int]     ,       __UpperCamelCase			:Callable[..., Tuple[torch.Tensor]]     ,       **__UpperCamelCase			:str		):
												super().__init__(**__UpperCamelCase		)
												A		=  args_dim
												A		=  nn.ModuleList([nn.Linear(__UpperCamelCase     ,       __UpperCamelCase		) for dim in args_dim.values()]		)
												A		=  domain_map


					def 	lowerCamelCase       (					self			:int     ,       __UpperCamelCase			:torch.Tensor		):
												A		=  [proj(__UpperCamelCase		) for proj in self.proj]

												return self.domain_map(*__UpperCamelCase		)







class 					_UpperCAmelCase   (	nn.Module ):


					def __init__(					self			:Dict     ,       __UpperCamelCase			:int		):
												super().__init__()
												A		=  function


					def 	lowerCamelCase       (					self			:List[str]     ,       __UpperCamelCase			:Any     ,       *__UpperCamelCase			:Any		):
												return self.function(__UpperCamelCase     ,       *__UpperCamelCase		)







class 					_UpperCAmelCase   :
					UpperCamelCase      =							42
					UpperCamelCase      =							42
					UpperCamelCase      =							42


					def __init__(					self			:Any     ,       __UpperCamelCase			:int = 1		):
												A		=  dim
												A		=  {k: dim * self.args_dim[k] for k in self.args_dim}


					def 	lowerCamelCase       (					self			:List[Any]     ,       __UpperCamelCase			:Dict		):
												if self.dim == 1:
																			return self.distribution_class(*__UpperCamelCase		)
												else:
																			return Independent(self.distribution_class(*__UpperCamelCase		)     ,       1		)


					def 	lowerCamelCase       (					self			:int     ,       __UpperCamelCase			:List[str]     ,       __UpperCamelCase			:Optional[torch.Tensor] = None     ,       __UpperCamelCase			:Optional[torch.Tensor] = None     ,       ):
												A		=  self._base_distribution(__UpperCamelCase		)
												if loc is None and scale is None:
																			return distr
												else:
																			return AffineTransformed(__UpperCamelCase     ,       loc=__UpperCamelCase     ,       scale=__UpperCamelCase     ,       event_dim=self.event_dim		)


					@property
					def 	lowerCamelCase       (					self			:List[Any]		):
												return () if self.dim == 1 else (self.dim,)


					@property
					def 	lowerCamelCase       (					self			:Tuple		):
												return len(self.event_shape		)


					@property
					def 	lowerCamelCase       (					self			:int		):
												return 0.0


					def 	lowerCamelCase       (					self			:str     ,       __UpperCamelCase			:int		):
												return ParameterProjection(
												    in_features=__UpperCamelCase     ,       args_dim=self.args_dim     ,       domain_map=LambdaLayer(self.domain_map		)     ,       )


					def 	lowerCamelCase       (					self			:List[Any]     ,       *__UpperCamelCase			:torch.Tensor		):
												raise NotImplementedError()


					@staticmethod
					def 	lowerCamelCase       (					__UpperCamelCase			:torch.Tensor		):
												return (x + torch.sqrt(torch.square(__UpperCamelCase		) + 4.0		)) / 2.0







class 					_UpperCAmelCase   (	lowercase_ ):
					UpperCamelCase      =							{"df": 1, "loc": 1, "scale": 1}
					UpperCamelCase      =							StudentT


					@classmethod
					def 	lowerCamelCase       (					cls			:List[str]     ,       __UpperCamelCase			:torch.Tensor     ,       __UpperCamelCase			:torch.Tensor     ,       __UpperCamelCase			:torch.Tensor		):
												A		=  cls.squareplus(__UpperCamelCase		).clamp_min(torch.finfo(scale.dtype		).eps		)
												A		=  2.0 + cls.squareplus(__UpperCamelCase		)
												return df.squeeze(-1		), loc.squeeze(-1		), scale.squeeze(-1		)







class 					_UpperCAmelCase   (	lowercase_ ):
					UpperCamelCase      =							{"loc": 1, "scale": 1}
					UpperCamelCase      =							Normal


					@classmethod
					def 	lowerCamelCase       (					cls			:List[Any]     ,       __UpperCamelCase			:torch.Tensor     ,       __UpperCamelCase			:torch.Tensor		):
												A		=  cls.squareplus(__UpperCamelCase		).clamp_min(torch.finfo(scale.dtype		).eps		)
												return loc.squeeze(-1		), scale.squeeze(-1		)







class 					_UpperCAmelCase   (	lowercase_ ):
					UpperCamelCase      =							{"total_count": 1, "logits": 1}
					UpperCamelCase      =							NegativeBinomial


					@classmethod
					def 	lowerCamelCase       (					cls			:str     ,       __UpperCamelCase			:torch.Tensor     ,       __UpperCamelCase			:torch.Tensor		):
												A		=  cls.squareplus(__UpperCamelCase		)
												return total_count.squeeze(-1		), logits.squeeze(-1		)


					def 	lowerCamelCase       (					self			:Tuple     ,       __UpperCamelCase			:List[str]		):
												A,	A		=  distr_args
												if self.dim == 1:
																			return self.distribution_class(total_count=__UpperCamelCase     ,       logits=__UpperCamelCase		)
												else:
																			return Independent(self.distribution_class(total_count=__UpperCamelCase     ,       logits=__UpperCamelCase		)     ,       1		)


					def 	lowerCamelCase       (					self			:List[str]     ,       __UpperCamelCase			:str     ,       __UpperCamelCase			:Optional[torch.Tensor] = None     ,       __UpperCamelCase			:Optional[torch.Tensor] = None		):
												A,	A		=  distr_args

												if scale is not None:
																			# See scaling property of Gamma.
																			logits += scale.log()

												return self._base_distribution((total_count, logits)		)


