






"""simple docstring"""



from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCAmelCase	:	Dict				=      logging.get_logger(__name__)

# TODO Update this
_lowerCAmelCase	:	Tuple				=      {
    'facebook/esm-1b': 'https://huggingface.co/facebook/esm-1b/resolve/main/config.json',
    # See all ESM models at https://huggingface.co/models?filter=esm
}




class 			lowerCAmelCase__			(    __snake_case   ):
   SCREAMING_SNAKE_CASE_    ='''esm'''
   def __init__(		self      :	Dict							,			snake_case__      :	Tuple=None							,			snake_case__      :	Any=None							,			snake_case__      :	Optional[Any]=None							,			snake_case__      :	List[Any]=7_6_8							,			snake_case__      :	Dict=1_2							,			snake_case__      :	str=1_2							,			snake_case__      :	Any=3_0_7_2							,			snake_case__      :	Tuple=0.1							,			snake_case__      :	int=0.1							,			snake_case__      :	List[str]=1_0_2_6							,			snake_case__      :	Any=0.02							,			snake_case__      :	Dict=1e-12							,			snake_case__      :	Union[str, Any]="absolute"							,			snake_case__      :	Any=True							,			snake_case__      :	Optional[int]=None							,			snake_case__      :	Tuple=False							,			snake_case__      :	str=False							,			snake_case__      :	List[Any]=None							,			snake_case__      :	List[Any]=None							,			**snake_case__      :	int							,			):


        '''simple docstring'''
        super().__init__(pad_token_id=snake_case__							,			mask_token_id=snake_case__							,			**snake_case__							)

        UpperCAmelCase__						:    Optional[Any]							=							vocab_size
        UpperCAmelCase__						:    int							=							hidden_size
        UpperCAmelCase__						:    Optional[Any]							=							num_hidden_layers
        UpperCAmelCase__						:    Optional[Any]							=							num_attention_heads
        UpperCAmelCase__						:    Dict							=							intermediate_size
        UpperCAmelCase__						:    int							=							hidden_dropout_prob
        UpperCAmelCase__						:    List[str]							=							attention_probs_dropout_prob
        UpperCAmelCase__						:    Dict							=							max_position_embeddings
        UpperCAmelCase__						:    List[str]							=							initializer_range
        UpperCAmelCase__						:    List[str]							=							layer_norm_eps
        UpperCAmelCase__						:    List[str]							=							position_embedding_type
        UpperCAmelCase__						:    Optional[Any]							=							use_cache
        UpperCAmelCase__						:    Dict							=							emb_layer_norm_before
        UpperCAmelCase__						:    List[Any]							=							token_dropout
        UpperCAmelCase__						:    List[Any]							=							is_folding_model
        if is_folding_model:
             if esmfold_config is None:
                  logger.info("No esmfold_config supplied for folding model, using default values."							)
                  UpperCAmelCase__						:    List[Any]							=							EsmFoldConfig()
             elif isinstance(snake_case__							,			snake_case__							):
                  UpperCAmelCase__						:    Optional[Any]							=							EsmFoldConfig(**snake_case__							)
             UpperCAmelCase__						:    int							=							esmfold_config
             if vocab_list is None:
                  logger.warning("No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"							)
                  UpperCAmelCase__						:    Tuple							=							get_default_vocab_list()
             else:
                  UpperCAmelCase__						:    Optional[int]							=							vocab_list
        else:
             UpperCAmelCase__						:    Dict							=							None
             UpperCAmelCase__						:    Tuple							=							None
        if self.esmfold_config is not None and getattr(self.esmfold_config							,			"use_esm_attn_map"							,			snake_case__							):
             raise ValueError("The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"							)




   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    int							=							super().to_dict()
        if isinstance(self.esmfold_config							,			snake_case__							):
             UpperCAmelCase__						:    List[str]							=							self.esmfold_config.to_dict()
        return output




@dataclass
class 			lowerCAmelCase__			:
   SCREAMING_SNAKE_CASE_    =None
   SCREAMING_SNAKE_CASE_    =True
   SCREAMING_SNAKE_CASE_    =False
   SCREAMING_SNAKE_CASE_    =False
   SCREAMING_SNAKE_CASE_    =False
   SCREAMING_SNAKE_CASE_    =0

   SCREAMING_SNAKE_CASE_    =True
   SCREAMING_SNAKE_CASE_    =False

   SCREAMING_SNAKE_CASE_    =128
   SCREAMING_SNAKE_CASE_    =None
   def       __a							(		self      :	Any							):


        '''simple docstring'''






        if self.trunk is None:
             UpperCAmelCase__						:    Tuple							=							TrunkConfig()
        elif isinstance(self.trunk							,			snake_case__							):
             UpperCAmelCase__						:    int							=							TrunkConfig(**self.trunk							)




   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							asdict(self							)
        UpperCAmelCase__						:    List[Any]							=							self.trunk.to_dict()
        return output




@dataclass
class 			lowerCAmelCase__			:
   SCREAMING_SNAKE_CASE_    =48
   SCREAMING_SNAKE_CASE_    =1024
   SCREAMING_SNAKE_CASE_    =128
   SCREAMING_SNAKE_CASE_    =32
   SCREAMING_SNAKE_CASE_    =32
   SCREAMING_SNAKE_CASE_    =32
   SCREAMING_SNAKE_CASE_    =0
   SCREAMING_SNAKE_CASE_    =0
   SCREAMING_SNAKE_CASE_    =False
   SCREAMING_SNAKE_CASE_    =4
   SCREAMING_SNAKE_CASE_    =128
   SCREAMING_SNAKE_CASE_    =None
   def       __a							(		self      :	Tuple							):


        '''simple docstring'''






        if self.structure_module is None:
             UpperCAmelCase__						:    Dict							=							StructureModuleConfig()
        elif isinstance(self.structure_module							,			snake_case__							):
             UpperCAmelCase__						:    Optional[int]							=							StructureModuleConfig(**self.structure_module							)

        if self.max_recycles <= 0:
             raise ValueError(f'`max_recycles` should be positive, got {self.max_recycles}.'							)
        if self.sequence_state_dim % self.sequence_state_dim != 0:
             raise ValueError(
                 "`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"
                 f' {self.sequence_state_dim} and {self.sequence_state_dim}.'							)
        if self.pairwise_state_dim % self.pairwise_state_dim != 0:
             raise ValueError(
                 "`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"
                 f' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'							)

        UpperCAmelCase__						:    Dict							=							self.sequence_state_dim // self.sequence_head_width
        UpperCAmelCase__						:    Optional[int]							=							self.pairwise_state_dim // self.pairwise_head_width

        if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
             raise ValueError(
                 "`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"
                 f' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'							)
        if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
             raise ValueError(
                 "`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"
                 f' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'							)
        if self.pairwise_state_dim % 2 != 0:
             raise ValueError(f'`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'							)

        if self.dropout >= 0.4:
             raise ValueError(f'`dropout` should not be greater than 0.4, got {self.dropout}.'							)




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


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							asdict(self							)
        UpperCAmelCase__						:    Dict							=							self.structure_module.to_dict()
        return output




@dataclass
class 			lowerCAmelCase__			:
   SCREAMING_SNAKE_CASE_    =384
   SCREAMING_SNAKE_CASE_    =128
   SCREAMING_SNAKE_CASE_    =16
   SCREAMING_SNAKE_CASE_    =128
   SCREAMING_SNAKE_CASE_    =12
   SCREAMING_SNAKE_CASE_    =4
   SCREAMING_SNAKE_CASE_    =8
   SCREAMING_SNAKE_CASE_    =0.1
   SCREAMING_SNAKE_CASE_    =8
   SCREAMING_SNAKE_CASE_    =1
   SCREAMING_SNAKE_CASE_    =2
   SCREAMING_SNAKE_CASE_    =7
   SCREAMING_SNAKE_CASE_    =10
   SCREAMING_SNAKE_CASE_    =1e-8
   SCREAMING_SNAKE_CASE_    =1e5
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        return asdict(self							)




def 	SCREAMING_SNAKE_CASE__       (  )->		int:






     '''simple docstring'''



     return (
         "<cls>",
         "<pad>",
         "<eos>",
         "<unk>",
         "L",
         "A",
         "G",
         "V",
         "S",
         "E",
         "R",
         "T",
         "I",
         "D",
         "P",
         "K",
         "Q",
         "N",
         "F",
         "Y",
         "M",
         "H",
         "W",
         "C",
         "X",
         "B",
         "U",
         "Z",
         "O",
         ".",
         "-",
         "<null_1>",
         "<mask>",
     )




'''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_							(  snake_case__		):




    '''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 					A							(				nn.Module		):





           def __init__(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )					->							Tuple:






               """simple docstring"""






               super().__init__()
               A			:  List[Any]   =		module
               A			:  Any   =		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			:  List[str]   =		(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 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     *SCREAMING_SNAKE_CASE		,     **SCREAMING_SNAKE_CASE      )					->							str:






               """simple docstring"""






               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 					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
     __magic_name__          =						'''bigscience/bloom-1b7'''

     # Constant values
     __magic_name__          =						2.1_09_65_95_52_69_25_74

     __magic_name__          =						'''Hello my name is'''
     __magic_name__          =						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'''		)
     __magic_name__          =						10





     def 							__lowerCAmelCase				(				self      )					->							str:






         """simple docstring"""






         A			:  Any   =		AutoTokenizer.from_pretrained(self.model_name      )





class 					A							(				__snake_case		):





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






         """simple docstring"""






         super().setUp()

         # Models and tokenizer
         A			:  Optional[Any]   =		AutoModelForCausalLM.from_pretrained(
             self.model_name		,     torch_dtype=torch.floataa		,     device_map='''auto'''      )
         A			:  List[Any]   =		AutoModelForCausalLM.from_pretrained(self.model_name		,     load_in_abit=SCREAMING_SNAKE_CASE		,     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      )					->							Optional[Any]:






         """simple docstring"""






         A			:  Any   =		self.model_abit.config

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

         A			:  List[str]   =		config.to_dict()
         A			:  List[Any]   =		config.to_diff_dict()

         A			:  Union[str, Any]   =		config.to_json_string()





     def 							__lowerCAmelCase				(				self      )					->							int:






         """simple docstring"""






         from bitsandbytes.nn import Paramsabit

         A			:  Optional[Any]   =		self.model_fpaa.get_memory_footprint()
         A			:  Any   =		self.model_abit.get_memory_footprint()

         self.assertAlmostEqual(mem_fpaa / mem_abit		,     self.EXPECTED_RELATIVE_DIFFERENCE      )
         A			:  Union[str, Any]   =		get_some_linear_layer(self.model_abit      )
         self.assertTrue(linear.weight.__class__ == Paramsabit      )





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """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(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 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A			:  List[Any]   =		self.tokenizer(self.input_text		,     return_tensors='''pt'''      )
         A			:  List[str]   =		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 							__lowerCAmelCase				(				self      )					->							List[str]:






         """simple docstring"""






         A			:  str   =		BitsAndBytesConfig()
         A			:  Optional[Any]   =		True

         A			:  int   =		AutoModelForCausalLM.from_pretrained(
             self.model_name		,     quantization_config=SCREAMING_SNAKE_CASE		,     device_map='''auto'''      )

         A			:  Union[str, Any]   =		self.tokenizer(self.input_text		,     return_tensors='''pt'''      )
         A			:  Tuple   =		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 							__lowerCAmelCase				(				self      )					->							Union[str, Any]:






         """simple docstring"""






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





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A			:  Union[str, Any]   =		BitsAndBytesConfig()

         with self.assertRaises(SCREAMING_SNAKE_CASE      ):
             A			:  int   =		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 							__lowerCAmelCase				(				self      )					->							Optional[Any]:






         """simple docstring"""






         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			:  int   =		self.tokenizer(self.input_text		,     return_tensors='''pt'''      )

         A			:  Dict   =		self.model_fpaa.to(torch.floataa      )
         A			:  Union[str, Any]   =		self.model_fpaa.generate(input_ids=encoded_input['''input_ids'''].to(0      )		,     max_new_tokens=10      )

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

         # Check this does not throw an error
         A			:  List[str]   =		self.model_fpaa.half()

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





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






         """simple docstring"""






         A			:  Tuple   =		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 					A							(				unittest.TestCase		):





     @classmethod
     def 							__lowerCAmelCase				(				cls      )					->							Union[str, Any]:






         """simple docstring"""






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





     def 							__lowerCAmelCase				(				self      )					->							str:






         """simple docstring"""






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





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






         """simple docstring"""






         from transformers import TaForConditionalGeneration

         A			:  int   =		TaForConditionalGeneration._keep_in_fpaa_modules
         A			:  Optional[Any]   =		None

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

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





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






         """simple docstring"""






         import bitsandbytes as bnb

         from transformers import TaForConditionalGeneration

         # test with `t5-small`
         A			:  List[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			:  List[Any]   =		self.tokenizer(self.input_text		,     return_tensors='''pt'''      ).to(0      )
         A			:  Any   =		model.generate(**SCREAMING_SNAKE_CASE      )

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





class 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         super().setUp()
         # model_name
         A			:  Tuple   =		'''bigscience/bloom-560m'''
         A			:  int   =		'''t5-small'''

         # Different types of model

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





     def 							__lowerCAmelCase				(				self      )					->							Tuple:






         """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      )					->							List[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 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         super().setUp()





     def 							__lowerCAmelCase				(				self      )					->							Tuple:






         """simple docstring"""






         del self.pipe

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





     def 							__lowerCAmelCase				(				self      )					->							Tuple:






         """simple docstring"""






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





@require_torch_multi_gpu
class 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self      )					->							Dict:






         """simple docstring"""






         super().setUp()





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






         """simple docstring"""






         A			:  Tuple   =		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			:  Optional[int]   =		self.tokenizer(self.input_text		,     return_tensors='''pt'''      )

         # Second real batch
         A			:  Optional[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 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self      )					->							str:






         """simple docstring"""






         A			:  int   =		'''facebook/opt-350m'''
         super().setUp()





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






         """simple docstring"""






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

         # Step 1: freeze all parameters
         A			:  Tuple   =		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[int]   =		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			:  Tuple   =		LoRALayer(module.q_proj		,     rank=16      )
                 A			:  Tuple   =		LoRALayer(module.k_proj		,     rank=16      )
                 A			:  str   =		LoRALayer(module.v_proj		,     rank=16      )

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

         # Step 4: Check if the gradient is not None
         with torch.cuda.amp.autocast():
             A			:  str   =		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 					A							(				__snake_case		):
     __magic_name__          =						'''gpt2-xl'''
     __magic_name__          =						3.31_91_85_48_54_15_21_87




