


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 		_A		(		_lowercase )					-> List[str]:


     """simple docstring"""


     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						):







           def __init__(			self: Union[str, Any],A_: nn.Module,A_: int					):






                '''simple docstring'''




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







           def 		snake_case_	(			self: Optional[Any],A_: int,*A_: Union[str, Any],**A_: str					):






                '''simple docstring'''




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



@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 	__lowerCamelCase			(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
       _lowercase											=     """bigscience/bloom-1b7"""

       # Constant values
       _lowercase											=     2.109659552692574

       _lowercase											=     """Hello my name is"""
       _lowercase											=     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"""						)
       _lowercase											=     10







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






            '''simple docstring'''




            __UpperCamelCase =			AutoTokenizer.from_pretrained(self.model_name					)



class 	__lowerCamelCase			(_a						):







       def 		snake_case_	(			self: int					):






            '''simple docstring'''




            super().setUp()

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







       def 		snake_case_	(			self: Optional[Any]					):






            '''simple docstring'''




            del self.model_fpaa
            del self.model_abit

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







       def 		snake_case_	(			self: str					):






            '''simple docstring'''




            __UpperCamelCase =			self.model_abit.config

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

            __UpperCamelCase =			config.to_dict()
            __UpperCamelCase =			config.to_diff_dict()

            __UpperCamelCase =			config.to_json_string()







       def 		snake_case_	(			self: Optional[int]					):






            '''simple docstring'''




            from bitsandbytes.nn import Paramsabit

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

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







       def 		snake_case_	(			self: str					):






            '''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(A_,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 		snake_case_	(			self: str					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: str					):






            '''simple docstring'''




            __UpperCamelCase =			BitsAndBytesConfig()
            __UpperCamelCase =			True

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

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







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: Dict					):






            '''simple docstring'''




            __UpperCamelCase =			BitsAndBytesConfig()

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







       def 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




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

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

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

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

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

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

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

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

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

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







       def 		snake_case_	(			self: Optional[int]					):






            '''simple docstring'''




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







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






            '''simple docstring'''




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







       def 		snake_case_	(			self: Dict					):






            '''simple docstring'''




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







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




            from transformers import TaForConditionalGeneration

            __UpperCamelCase =			TaForConditionalGeneration._keep_in_fpaa_modules
            __UpperCamelCase =			None

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

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







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






            '''simple docstring'''




            import bitsandbytes as bnb

            from transformers import TaForConditionalGeneration

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

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

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



class 	__lowerCamelCase			(_a						):







       def 		snake_case_	(			self: Tuple					):






            '''simple docstring'''




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

            # Different types of model

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







       def 		snake_case_	(			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 		snake_case_	(			self: 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 	__lowerCamelCase			(_a						):







       def 		snake_case_	(			self: Any					):






            '''simple docstring'''




            super().setUp()







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




            del self.pipe

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







       def 		snake_case_	(			self: str					):






            '''simple docstring'''




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



@require_torch_multi_gpu
class 	__lowerCamelCase			(_a						):







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






            '''simple docstring'''




            super().setUp()







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




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

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



class 	__lowerCamelCase			(_a						):







       def 		snake_case_	(			self: List[str]					):






            '''simple docstring'''




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







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






            '''simple docstring'''




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

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

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

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

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

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

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

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



class 	__lowerCamelCase			(_a						):
       _lowercase											=     """gpt2-xl"""
       _lowercase											=     3.3191854854152187










"""simple docstring"""
class _lowerCAmelCase						:







						def __init__(    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_=None   ,				lowerCAmelCase_=None							)		->							int:
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    data
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    previous
										_SCREAMING_SNAKE_CASE   :		Any       		=    next_node







						def __str__(    self							)		->							str:
										return F"""{self.data}"""







						def 			A      (    self							)		->							int:
										return self.data







						def 			A      (    self							)		->							Dict:
										return self.next







						def 			A      (    self							)		->							Union[str, Any]:
										return self.previous







class _lowerCAmelCase						:







						def __init__(    self   ,				lowerCAmelCase_							)		->							List[str]:
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    head







						def __iter__(    self							)		->							Any:
										return self







						def 			A      (    self							)		->							Dict:
										if not self.current:
														raise StopIteration
										else:
														_SCREAMING_SNAKE_CASE   :		Any       		=    self.current.get_data()
														_SCREAMING_SNAKE_CASE   :		Dict       		=    self.current.get_next()
														return value







class _lowerCAmelCase						:







						def __init__(    self							)		->							Dict:
										_SCREAMING_SNAKE_CASE   :		Dict       		=    None  # First node in list
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    None  # Last node in list







						def __str__(    self							)		->							Any:
										_SCREAMING_SNAKE_CASE   :		str       		=    self.head
										_SCREAMING_SNAKE_CASE   :		int       		=    []
										while current is not None:
														nodes.append(current.get_data()							)
														_SCREAMING_SNAKE_CASE   :		List[Any]       		=    current.get_next()
										return " ".join(str(lowerCAmelCase_							) for node in nodes							)







						def __contains__(    self   ,				lowerCAmelCase_							)		->							Any:
										_SCREAMING_SNAKE_CASE   :		Any       		=    self.head
										while current:
														if current.get_data() == value:
																		return True
														_SCREAMING_SNAKE_CASE   :		int       		=    current.get_next()
										return False







						def __iter__(    self							)		->							str:
										return LinkedListIterator(self.head							)







						def 			A      (    self							)		->							Dict:
										if self.head:
														return self.head.get_data()
										return None







						def 			A      (    self							)		->							List[Any]:
										if self.tail:
														return self.tail.get_data()
										return None







						def 			A      (    self   ,				lowerCAmelCase_							)		->							None:
										if self.head is None:
														_SCREAMING_SNAKE_CASE   :		List[Any]       		=    node
														_SCREAMING_SNAKE_CASE   :		Dict       		=    node
										else:
														self.insert_before_node(self.head   ,				lowerCAmelCase_							)







						def 			A      (    self   ,				lowerCAmelCase_							)		->							None:
										if self.head is None:
														self.set_head(lowerCAmelCase_							)
										else:
														self.insert_after_node(self.tail   ,				lowerCAmelCase_							)







						def 			A      (    self   ,				lowerCAmelCase_							)		->							None:
										_SCREAMING_SNAKE_CASE   :		Any       		=    Node(lowerCAmelCase_							)
										if self.head is None:
														self.set_head(lowerCAmelCase_							)
										else:
														self.set_tail(lowerCAmelCase_							)







						def 			A      (    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_							)		->							None:
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    node
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    node.previous

										if node.get_previous() is None:
														_SCREAMING_SNAKE_CASE   :		List[Any]       		=    node_to_insert
										else:
														_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    node_to_insert

										_SCREAMING_SNAKE_CASE   :		Any       		=    node_to_insert







						def 			A      (    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_							)		->							None:
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    node
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    node.next

										if node.get_next() is None:
														_SCREAMING_SNAKE_CASE   :		Tuple       		=    node_to_insert
										else:
														_SCREAMING_SNAKE_CASE   :		Any       		=    node_to_insert

										_SCREAMING_SNAKE_CASE   :		List[str]       		=    node_to_insert







						def 			A      (    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_							)		->							None:
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    1
										_SCREAMING_SNAKE_CASE   :		str       		=    Node(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    self.head
										while node:
														if current_position == position:
																		self.insert_before_node(lowerCAmelCase_   ,				lowerCAmelCase_							)
																		return
														current_position += 1
														_SCREAMING_SNAKE_CASE   :		List[str]       		=    node.next
										self.insert_after_node(self.tail   ,				lowerCAmelCase_							)







						def 			A      (    self   ,				lowerCAmelCase_							)		->							Node:
										_SCREAMING_SNAKE_CASE   :		Tuple       		=    self.head
										while node:
														if node.get_data() == item:
																		return node
														_SCREAMING_SNAKE_CASE   :		Dict       		=    node.get_next()
										raise Exception('Node not found'							)







						def 			A      (    self   ,				lowerCAmelCase_							)		->							int:
										if (node := self.get_node(lowerCAmelCase_							)) is not None:
														if node == self.head:
																		_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    self.head.get_next()

														if node == self.tail:
																		_SCREAMING_SNAKE_CASE   :		int       		=    self.tail.get_previous()

														self.remove_node_pointers(lowerCAmelCase_							)







						@staticmethod
						def 			A      (    lowerCAmelCase_							)		->							None:
										if node.get_next():
														_SCREAMING_SNAKE_CASE   :		List[Any]       		=    node.previous

										if node.get_previous():
														_SCREAMING_SNAKE_CASE   :		str       		=    node.next

										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    None
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    None







						def 			A      (    self							)		->							List[str]:
										return self.head is None







def 				lowercase__		(					):
				pass


if __name__ == "__main__":
			import doctest

			doctest.testmod()

