



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     __lowercase(	UpperCAmelCase__ ):


      """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):
          """simple docstring"""



          def __init__(self				,			__a				,			__a    ):






                '''simple docstring'''



                super().__init__()
                lowerCamelCase								=		module
                lowerCamelCase								=		nn.Sequential(
                    nn.Linear(module.in_features				,			__a				,			bias=__a    )				,			nn.Linear(__a				,			module.out_features				,			bias=__a    )				,			)
                lowerCamelCase								=		(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				,			__a				,			*__a				,			**__a    ):






                '''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):
     """simple docstring"""



     _A       =			'bigscience/bloom-1b7'

     # Constant values
     _A       =			2.1_09_65_95_52_69_25_74

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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		AutoTokenizer.from_pretrained(self.model_name    )





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



     def 						_a      (self    ):






           '''simple docstring'''



           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=__a				,			device_map="auto"    )







     def 						_a      (self    ):






           '''simple docstring'''



           del self.model_fpaa
           del self.model_abit

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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_abit.config

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

           lowerCamelCase								=		config.to_dict()
           lowerCamelCase								=		config.to_diff_dict()

           lowerCamelCase								=		config.to_json_string()







     def 						_a      (self    ):






           '''simple docstring'''



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






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






           '''simple docstring'''



           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=__a    )				,			self.EXPECTED_OUTPUTS    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		BitsAndBytesConfig()
           lowerCamelCase								=		True

           lowerCamelCase								=		AutoModelForCausalLM.from_pretrained(
               self.model_name				,			quantization_config=__a				,			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=__a    )				,			self.EXPECTED_OUTPUTS    )







     def 						_a      (self    ):






           '''simple docstring'''



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







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		BitsAndBytesConfig()

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







     def 						_a      (self    ):






           '''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
           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 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		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):
     """simple docstring"""



     @classmethod
     def 						_a      (cls    ):






           '''simple docstring'''



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






           '''simple docstring'''



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







     def 						_a      (self    ):






           '''simple docstring'''



           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=__a				,			device_map="auto"    )
           lowerCamelCase								=		self.tokenizer(self.input_text				,			return_tensors="pt"    ).to(0    )
           lowerCamelCase								=		model.generate(**__a    )

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







     def 						_a      (self    ):






           '''simple docstring'''



           import bitsandbytes as bnb

           from transformers import TaForConditionalGeneration

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

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

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





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



     def 						_a      (self    ):






           '''simple docstring'''



           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=__a				,			device_map="auto"    )
           # Sequence classification model
           lowerCamelCase								=		AutoModelForSequenceClassification.from_pretrained(
               self.model_name				,			load_in_abit=__a				,			device_map="auto"    )
           # CausalLM model
           lowerCamelCase								=		AutoModelForCausalLM.from_pretrained(self.model_name				,			load_in_abit=__a				,			device_map="auto"    )
           # Seq2seq model
           lowerCamelCase								=		AutoModelForSeqaSeqLM.from_pretrained(
               self.seq_to_seq_name				,			load_in_abit=__a				,			device_map="auto"    )







     def 						_a      (self    ):






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






           '''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__  (						UpperCAmelCase_):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           super().setUp()







     def 						_a      (self    ):






           '''simple docstring'''



           del self.pipe

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







     def 						_a      (self    ):






           '''simple docstring'''



           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__  (						UpperCAmelCase_):
     """simple docstring"""



     def 						_a      (self    ):






           '''simple docstring'''



           super().setUp()







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		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
           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=__a    )				,			self.EXPECTED_OUTPUTS    )





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



     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		"facebook/opt-350m"
           super().setUp()







     def 						_a      (self    ):






           '''simple docstring'''



           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=__a    )

           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(__a    )    ):
                       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(**__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__  (						UpperCAmelCase_):
     """simple docstring"""



     _A       =			'gpt2-xl'
     _A       =			3.31_91_85_48_54_15_21_87




import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding





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



     _A       =			['image_processor', 'tokenizer']
     _A       =			'CLIPImageProcessor'
     _A       =			('CLIPTokenizer', 'CLIPTokenizerFast')







     def __init__(self				,			__a=None				,			__a=None				,			**__a    ):






           '''simple docstring'''



           lowerCamelCase								=		None
           if "feature_extractor" in kwargs:
                 warnings.warn(
                     "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
                     " instead."				,			__a				,			)
                 lowerCamelCase								=		kwargs.pop("feature_extractor"    )

           lowerCamelCase								=		image_processor if image_processor is not None else feature_extractor
           if image_processor is None:
                 raise ValueError("You need to specify an `image_processor`."    )
           if tokenizer is None:
                 raise ValueError("You need to specify a `tokenizer`."    )

           super().__init__(__a				,			__a    )







     def __call__(self				,			__a=None				,			__a=None				,			__a=None				,			**__a    ):






           '''simple docstring'''



           if text is None and images is None:
                 raise ValueError("You have to specify either text or images. Both cannot be none."    )

           if text is not None:
                 lowerCamelCase								=		self.tokenizer(__a				,			return_tensors=__a				,			**__a    )

           if images is not None:
                 lowerCamelCase								=		self.image_processor(__a				,			return_tensors=__a				,			**__a    )

           if text is not None and images is not None:
                 lowerCamelCase								=		image_features.pixel_values
                 return encoding
           elif text is not None:
                 return encoding
           else:
                 return BatchEncoding(data=dict(**__a    )				,			tensor_type=__a    )







     def 						_a      (self				,			*__a				,			**__a    ):






           '''simple docstring'''



           return self.tokenizer.batch_decode(*__a				,			**__a    )







     def 						_a      (self				,			*__a				,			**__a    ):






           '''simple docstring'''



           return self.tokenizer.decode(*__a				,			**__a    )







     @property
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.tokenizer.model_input_names
           lowerCamelCase								=		self.image_processor.model_input_names
           return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names    )    )







     @property
     def 						_a      (self    ):






           '''simple docstring'''



           warnings.warn(
               "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."				,			__a				,			)
           return self.image_processor_class







     @property
     def 						_a      (self    ):






           '''simple docstring'''



           warnings.warn(
               "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."				,			__a				,			)
           return self.image_processor