
"""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       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 		UpperCAmelCase_						(nn.Module    ):
        """simple docstring"""




        def __init__(			self							:			Any  , a_							:			nn.Module  , a_							:			int			)->       Tuple:



             """simple docstring"""

             super().__init__()
             UpperCAmelCase_    :     Tuple            =	module
             UpperCAmelCase_    :     Optional[Any]            =	nn.Sequential(
                 nn.Linear(module.in_features  , a_  , bias=a_			)  , nn.Linear(a_  , module.out_features  , bias=a_			)  , )
             UpperCAmelCase_    :     Dict            =	(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 	a   (			self							:			int  , a_							:			List[str]  , *a_							:			Optional[int]  , **a_							:			List[Any]			)->       List[str]:



             """simple docstring"""

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



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




 UpperCamelCase_    :  Any					 =     """bigscience/bloom-1b7"""

 # Constant values
 UpperCamelCase_    :  Union[str, Any]					 =     2.1_09_65_95_52_69_25_74

 UpperCamelCase_    :  List[Any]					 =     """Hello my name is"""
 UpperCamelCase_    :  Dict					 =     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_    :  Union[str, Any]					 =     10




 def 	a   (			self							:			Tuple			)->       Tuple:



      """simple docstring"""

      UpperCAmelCase_    :     Union[str, Any]            =	AutoTokenizer.from_pretrained(self.model_name			)



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




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



      """simple docstring"""

      super().setUp()

      # Models and tokenizer
      UpperCAmelCase_    :     Any            =	AutoModelForCausalLM.from_pretrained(
          self.model_name  , torch_dtype=torch.floataa  , device_map="""auto"""			)
      UpperCAmelCase_    :     Tuple            =	AutoModelForCausalLM.from_pretrained(self.model_name  , load_in_abit=a_  , device_map="""auto"""			)




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



      """simple docstring"""

      del self.model_fpaa
      del self.model_abit

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




 def 	a   (			self							:			Dict			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	self.model_abit.config

      self.assertTrue(hasattr(a_  , """quantization_config"""			)			)

      UpperCAmelCase_    :     str            =	config.to_dict()
      UpperCAmelCase_    :     Tuple            =	config.to_diff_dict()

      UpperCAmelCase_    :     Dict            =	config.to_json_string()




 def 	a   (			self							:			str			)->       int:



      """simple docstring"""

      from bitsandbytes.nn import Paramsabit

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

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




 def 	a   (			self							:			Dict			)->       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(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 	a   (			self							:			Optional[Any]			)->       Dict:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			)
      UpperCAmelCase_    :     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=a_			)  , self.EXPECTED_OUTPUTS			)




 def 	a   (			self							:			int			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	BitsAndBytesConfig()
      UpperCAmelCase_    :     List[str]            =	True

      UpperCAmelCase_    :     Optional[int]            =	AutoModelForCausalLM.from_pretrained(
          self.model_name  , quantization_config=a_  , device_map="""auto"""			)

      UpperCAmelCase_    :     int            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			)
      UpperCAmelCase_    :     Union[str, Any]            =	model_abit_from_config.generate(
          input_ids=encoded_input["""input_ids"""].to(0			)  , max_new_tokens=10			)

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




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



      """simple docstring"""

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




 def 	a   (			self							:			Optional[int]			)->       Union[str, Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	BitsAndBytesConfig()

      with self.assertRaises(a_			):
           UpperCAmelCase_    :     Dict            =	AutoModelForCausalLM.from_pretrained(
               self.model_name  , quantization_config=a_  , load_in_abit=a_  , device_map="""auto"""  , bnb_abit_quant_type="""nf4"""  , )




 def 	a   (			self							:			str			)->       Any:



      """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_    :     Optional[Any]            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			)

      UpperCAmelCase_    :     Union[str, Any]            =	self.model_fpaa.to(torch.floataa			)
      UpperCAmelCase_    :     List[Any]            =	self.model_fpaa.generate(input_ids=encoded_input["""input_ids"""].to(0			)  , max_new_tokens=10			)

      # Check this does not throw an error
      UpperCAmelCase_    :     Union[str, Any]            =	self.model_fpaa.to("""cpu"""			)

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

      # Check this does not throw an error
      UpperCAmelCase_    :     Union[str, Any]            =	self.model_fpaa.float()






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



      """simple docstring"""

      UpperCAmelCase_    :     str            =	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 		UpperCAmelCase_						(unittest.TestCase    ):
 """simple docstring"""




 @classmethod
 def 	a   (			cls							:			Dict			)->       Optional[Any]:



      """simple docstring"""

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




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



      """simple docstring"""

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




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



      """simple docstring"""

      from transformers import TaForConditionalGeneration

      UpperCAmelCase_    :     str            =	TaForConditionalGeneration._keep_in_fpaa_modules
      UpperCAmelCase_    :     List[str]            =	None

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

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






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



      """simple docstring"""

      import bitsandbytes as bnb

      from transformers import TaForConditionalGeneration

      # test with `t5-small`
      UpperCAmelCase_    :     Optional[Any]            =	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_    :     str            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			).to(0			)
      UpperCAmelCase_    :     Optional[int]            =	model.generate(**a_			)

      # test with `flan-t5-small`
      UpperCAmelCase_    :     Dict            =	TaForConditionalGeneration.from_pretrained(
          self.dense_act_model_name  , load_in_abit=a_  , device_map="""auto"""			)
      UpperCAmelCase_    :     Dict            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			).to(0			)
      UpperCAmelCase_    :     Union[str, Any]            =	model.generate(**a_			)



class 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




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



      """simple docstring"""

      super().setUp()
      # model_name
      UpperCAmelCase_    :     Optional[int]            =	"""bigscience/bloom-560m"""
      UpperCAmelCase_    :     Any            =	"""t5-small"""

      # Different types of model

      UpperCAmelCase_    :     Optional[int]            =	AutoModel.from_pretrained(self.model_name  , load_in_abit=a_  , device_map="""auto"""			)
      # Sequence classification model
      UpperCAmelCase_    :     Any            =	AutoModelForSequenceClassification.from_pretrained(
          self.model_name  , load_in_abit=a_  , device_map="""auto"""			)
      # CausalLM model
      UpperCAmelCase_    :     Tuple            =	AutoModelForCausalLM.from_pretrained(self.model_name  , load_in_abit=a_  , device_map="""auto"""			)
      # Seq2seq model
      UpperCAmelCase_    :     int            =	AutoModelForSeqaSeqLM.from_pretrained(
          self.seq_to_seq_name  , load_in_abit=a_  , device_map="""auto"""			)




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



      """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 	a   (			self							:			List[str]			)->       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 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def 	a   (			self							:			Tuple			)->       Tuple:



      """simple docstring"""

      super().setUp()




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



      """simple docstring"""

      del self.pipe

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






 def 	a   (			self							:			int			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	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_    :     str            =	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 	a   (			self							:			str			)->       Union[str, Any]:



      """simple docstring"""

      super().setUp()






 def 	a   (			self							:			str			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	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_    :     str            =	self.tokenizer(self.input_text  , return_tensors="""pt"""			)

      # Second real batch
      UpperCAmelCase_    :     List[str]            =	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 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 def 	a   (			self							:			int			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	"""facebook/opt-350m"""
      super().setUp()






 def 	a   (			self							:			str			)->       Any:



      """simple docstring"""

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

      # Step 1: freeze all parameters
      UpperCAmelCase_    :     Tuple            =	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_    :     Union[str, Any]            =	False  # freeze the model - train adapters later
           if param.ndim == 1:
                # cast the small parameters (e.g. layernorm) to fp32 for stability
                UpperCAmelCase_    :     Optional[int]            =	param.data.to(torch.floataa			)

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

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

      # Step 4: Check if the gradient is not None
      with torch.cuda.amp.autocast():
           UpperCAmelCase_    :     Optional[Any]            =	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 		UpperCAmelCase_						(lowerCamelCase_    ):
 """simple docstring"""




 UpperCamelCase_    :  str					 =     """gpt2-xl"""
 UpperCamelCase_    :  int					 =     3.31_91_85_48_54_15_21_87



"""simple docstring"""







from . import __version__

# Backward compatibility imports, to make sure all those objects can be found in file_utils
from .utils import (
    CLOUDFRONT_DISTRIB_PREFIX,
    CONFIG_NAME,
    DISABLE_TELEMETRY,
    DUMMY_INPUTS,
    DUMMY_MASK,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    FEATURE_EXTRACTOR_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_PREFIX,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    MODEL_CARD_NAME,
    MULTIPLE_CHOICE_DUMMY_INPUTS,
    PYTORCH_PRETRAINED_BERT_CACHE,
    PYTORCH_TRANSFORMERS_CACHE,
    S3_BUCKET_PREFIX,
    SENTENCEPIECE_UNDERLINE,
    SPIECE_UNDERLINE,
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    TORCH_FX_REQUIRED_VERSION,
    TRANSFORMERS_CACHE,
    TRANSFORMERS_DYNAMIC_MODULE_NAME,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    ContextManagers,
    DummyObject,
    EntryNotFoundError,
    ExplicitEnum,
    ModelOutput,
    PaddingStrategy,
    PushToHubMixin,
    RepositoryNotFoundError,
    RevisionNotFoundError,
    TensorType,
    _LazyModule,
    add_code_sample_docstrings,
    add_end_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    cached_property,
    copy_func,
    default_cache_path,
    define_sagemaker_information,
    get_cached_models,
    get_file_from_repo,
    get_full_repo_name,
    get_torch_version,
    has_file,
    http_user_agent,
    is_apex_available,
    is_bsa_available,
    is_coloredlogs_available,
    is_datasets_available,
    is_detectrona_available,
    is_faiss_available,
    is_flax_available,
    is_ftfy_available,
    is_in_notebook,
    is_ipex_available,
    is_librosa_available,
    is_offline_mode,
    is_onnx_available,
    is_pandas_available,
    is_phonemizer_available,
    is_protobuf_available,
    is_psutil_available,
    is_pyanvml_available,
    is_pyctcdecode_available,
    is_pytesseract_available,
    is_pytorch_quantization_available,
    is_rjieba_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_scipy_available,
    is_sentencepiece_available,
    is_seqio_available,
    is_sklearn_available,
    is_soundfile_availble,
    is_spacy_available,
    is_speech_available,
    is_tensor,
    is_tensorflow_probability_available,
    is_tfaonnx_available,
    is_tf_available,
    is_timm_available,
    is_tokenizers_available,
    is_torch_available,
    is_torch_bfaa_available,
    is_torch_cuda_available,
    is_torch_fx_available,
    is_torch_fx_proxy,
    is_torch_mps_available,
    is_torch_tfaa_available,
    is_torch_tpu_available,
    is_torchaudio_available,
    is_training_run_on_sagemaker,
    is_vision_available,
    replace_return_docstrings,
    requires_backends,
    to_numpy,
    to_py_obj,
    torch_only_method,
)

