




import gc
import importlib.metadata
import tempfile
import unittest

from packaging import version

from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from transformers.testing_utils import (
    is_torch_available,
    require_accelerate,
    require_bitsandbytes,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)




def 					lowerCamelCase__       (						A__       :					Union[str, Any]	):





       '''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: List[str]    ,       UpperCamelCase_: nn.Module    ,       UpperCamelCase_: int   ):
                    super().__init__()
                    __lowerCamelCase   =						module
                    __lowerCamelCase   =						nn.Sequential(
                        nn.Linear(module.in_features    ,       UpperCamelCase_    ,       bias=UpperCamelCase_   )    ,       nn.Linear(UpperCamelCase_    ,       module.out_features    ,       bias=UpperCamelCase_   )    ,       )
                    __lowerCamelCase   =						(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: Union[str, Any]    ,       UpperCamelCase_: List[str]    ,       *UpperCamelCase_: str    ,       **UpperCamelCase_: List[str]   ):
                    return self.module(UpperCamelCase_    ,       *UpperCamelCase_    ,       **UpperCamelCase_   ) + self.adapter(UpperCamelCase_   )





@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
      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__      :	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__      :	Any    		=   10
      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             # Models and tokenizer
             __lowerCamelCase   =						AutoTokenizer.from_pretrained(self.model_name   )





class     lowerCamelCase__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: List[str]   ):
             super().setUp()

             # Models and tokenizer
             __lowerCamelCase   =						AutoModelForCausalLM.from_pretrained(
                 self.model_name    ,       torch_dtype=torch.floataa    ,       device_map="""auto"""   )
             __lowerCamelCase   =						AutoModelForCausalLM.from_pretrained(self.model_name    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""   )
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             del self.model_fpaa
             del self.model_abit

             gc.collect()
             torch.cuda.empty_cache()
      def 			lowerCAmelCase__						(				self: Any   ):
             __lowerCamelCase   =						self.model_abit.config

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

             __lowerCamelCase   =						config.to_dict()
             __lowerCamelCase   =						config.to_diff_dict()

             __lowerCamelCase   =						config.to_json_string()
      def 			lowerCAmelCase__						(				self: Dict   ):
             from bitsandbytes.nn import Paramsabit

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

             self.assertAlmostEqual(mem_fpaa / mem_abit    ,       self.EXPECTED_RELATIVE_DIFFERENCE   )
             __lowerCamelCase   =						get_some_linear_layer(self.model_abit   )
             self.assertTrue(linear.weight.__class__ == Paramsabit   )
      def 			lowerCAmelCase__						(				self: List[str]   ):
             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: Any   ):
             __lowerCamelCase   =						self.tokenizer(self.input_text    ,       return_tensors="""pt"""   )
             __lowerCamelCase   =						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: Union[str, Any]   ):
             __lowerCamelCase   =						BitsAndBytesConfig()
             __lowerCamelCase   =						True

             __lowerCamelCase   =						AutoModelForCausalLM.from_pretrained(
                 self.model_name    ,       quantization_config=UpperCamelCase_    ,       device_map="""auto"""   )

             __lowerCamelCase   =						self.tokenizer(self.input_text    ,       return_tensors="""pt"""   )
             __lowerCamelCase   =						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: Dict   ):
             with self.assertRaises(UpperCamelCase_   ), tempfile.TemporaryDirectory() as tmpdirname:
                    self.model_abit.save_pretrained(UpperCamelCase_   )
      def 			lowerCAmelCase__						(				self: Any   ):
             __lowerCamelCase   =						BitsAndBytesConfig()

             with self.assertRaises(UpperCamelCase_   ):
                    __lowerCamelCase   =						AutoModelForCausalLM.from_pretrained(
                        self.model_name    ,       quantization_config=UpperCamelCase_    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""    ,       bnb_abit_quant_type="""nf4"""    ,       )
      def 			lowerCAmelCase__						(				self: Dict   ):
             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
             __lowerCamelCase   =						self.tokenizer(self.input_text    ,       return_tensors="""pt"""   )

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

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

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

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




      def 			lowerCAmelCase__						(				self: Tuple   ):
             __lowerCamelCase   =						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     lowerCamelCase__(					unittest.TestCase):
      @classmethod
      def 			lowerCAmelCase__						(				cls: List[Any]   ):
             __lowerCamelCase   =						"""t5-small"""
             __lowerCamelCase   =						"""google/flan-t5-small"""  # flan-t5 uses dense-act instead of dense-relu-dense
             __lowerCamelCase   =						AutoTokenizer.from_pretrained(cls.model_name   )
             __lowerCamelCase   =						"""Translate in German: Hello, my dog is cute"""
      def 			lowerCAmelCase__						(				self: Dict   ):
             gc.collect()
             torch.cuda.empty_cache()
      def 			lowerCAmelCase__						(				self: Any   ):
             from transformers import TaForConditionalGeneration

             __lowerCamelCase   =						TaForConditionalGeneration._keep_in_fpaa_modules
             __lowerCamelCase   =						None

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

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




      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             import bitsandbytes as bnb

             from transformers import TaForConditionalGeneration

             # test with `t5-small`
             __lowerCamelCase   =						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   )   )

             __lowerCamelCase   =						self.tokenizer(self.input_text    ,       return_tensors="""pt"""   ).to(0   )
             __lowerCamelCase   =						model.generate(**UpperCamelCase_   )

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





class     lowerCamelCase__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: Dict   ):
             super().setUp()
             # model_name
             __lowerCamelCase   =						"""bigscience/bloom-560m"""
             __lowerCamelCase   =						"""t5-small"""

             # Different types of model

             __lowerCamelCase   =						AutoModel.from_pretrained(self.model_name    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""   )
             # Sequence classification model
             __lowerCamelCase   =						AutoModelForSequenceClassification.from_pretrained(
                 self.model_name    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""   )
             # CausalLM model
             __lowerCamelCase   =						AutoModelForCausalLM.from_pretrained(self.model_name    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""   )
             # Seq2seq model
             __lowerCamelCase   =						AutoModelForSeqaSeqLM.from_pretrained(
                 self.seq_to_seq_name    ,       load_in_abit=UpperCamelCase_    ,       device_map="""auto"""   )
      def 			lowerCAmelCase__						(				self: int   ):
             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]   ):
             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__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             super().setUp()
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             del self.pipe

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




      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             __lowerCamelCase   =						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
             __lowerCamelCase   =						self.pipe(self.input_text   )
             self.assertIn(pipeline_output[0]["""generated_text"""]    ,       self.EXPECTED_OUTPUTS   )





@require_torch_multi_gpu
class     lowerCamelCase__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: Tuple   ):
             super().setUp()




      def 			lowerCAmelCase__						(				self: Tuple   ):
             __lowerCamelCase   =						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
             __lowerCamelCase   =						self.tokenizer(self.input_text    ,       return_tensors="""pt"""   )

             # Second real batch
             __lowerCamelCase   =						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     lowerCamelCase__(					__lowerCamelCase):
      def 			lowerCAmelCase__						(				self: Tuple   ):
             __lowerCamelCase   =						"""facebook/opt-350m"""
             super().setUp()




      def 			lowerCAmelCase__						(				self: int   ):
             if version.parse(importlib.metadata.version("""bitsandbytes"""   )   ) < version.parse("""0.37.0"""   ):
                    return

             # Step 1: freeze all parameters
             __lowerCamelCase   =						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():
                    __lowerCamelCase   =						False  # freeze the model - train adapters later
                    if param.ndim == 1:
                           # cast the small parameters (e.g. layernorm) to fp32 for stability
                           __lowerCamelCase   =						param.data.to(torch.floataa   )

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

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

             # Step 4: Check if the gradient is not None
             with torch.cuda.amp.autocast():
                    __lowerCamelCase   =						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     lowerCamelCase__(					__lowerCamelCase):
      UpperCAmelCase__      :	Tuple    		=   'gpt2-xl'
      UpperCAmelCase__      :	Union[str, Any]    		=   3.31_91_85_48_54_15_21_87







from __future__ import annotations




def 					lowerCamelCase__       (						A__       :					list[int | float]  ,			A__       :					int  ,			A__       :					int	):





       '''simple docstring'''



       if len(A__	) == 0:
              raise ValueError("""find_max() arg is an empty sequence"""	)
       if (
           left >= len(A__	)
           or left < -len(A__	)
           or right >= len(A__	)
           or right < -len(A__	)
       ):
              raise IndexError("""list index out of range"""	)
       if left == right:
              return nums[left]
       __lowerCamelCase   =						(left + right) >> 1  # the middle
       __lowerCamelCase   =						find_max(A__  ,			A__  ,			A__	)  # find max in range[left, mid]
       __lowerCamelCase   =						find_max(A__  ,			mid + 1  ,			A__	)  # find max in range[mid + 1, right]

       return left_max if left_max >= right_max else right_max


if __name__ == "__main__":
       import doctest

       doctest.testmod(verbose=True)

