



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     (					_UpperCAmelCase    :       Dict   )    ->			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       :  Optional[int]			,		A       :  nn.Module			,		A       :  int)       ->   Optional[int]:



													"""simple docstring"""






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



													"""simple docstring"""






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







@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class       __lowerCAmelCase   (      unittest.TestCase						):
			UpperCamelCase        =					"""bigscience/bloom-1b7"""

			# Constant values
			UpperCamelCase        =					2.109_659_552_692_574

			UpperCamelCase        =					"""Hello my name is"""
			UpperCamelCase        =					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        =					1_0



			def    _lowerCamelCase	(						self       :  Any)       ->   Tuple:



						"""simple docstring"""






						_UpperCAmelCase			=       AutoTokenizer.from_pretrained(self.model_name)







class       __lowerCAmelCase   (      lowerCamelCase_						):



			def    _lowerCamelCase	(						self       :  Optional[Any])       ->   Dict:



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



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



						"""simple docstring"""






						del self.model_fpaa
						del self.model_abit

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



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



						"""simple docstring"""






						_UpperCAmelCase			=       self.model_abit.config

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

						_UpperCAmelCase			=       config.to_dict()
						_UpperCAmelCase			=       config.to_diff_dict()

						_UpperCAmelCase			=       config.to_json_string()



			def    _lowerCamelCase	(						self       :  List[str])       ->   Any:



						"""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    _lowerCamelCase	(						self       :  Optional[int])       ->   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(_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       :  Optional[Any])       ->   Any:



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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       BitsAndBytesConfig()
						_UpperCAmelCase			=       True

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



			def    _lowerCamelCase	(						self       :  List[str])       ->   Any:



						"""simple docstring"""






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



			def    _lowerCamelCase	(						self       :  str)       ->   Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase			=       BitsAndBytesConfig()

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



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



						"""simple docstring"""






						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
						_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    _lowerCamelCase	(						self       :  Any)       ->   Tuple:



						"""simple docstring"""






						_UpperCAmelCase			=       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       :  Optional[Any])       ->   int:



						"""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    _lowerCamelCase	(						self       :  Optional[int])       ->   Dict:



						"""simple docstring"""






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



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



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

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



			def    _lowerCamelCase	(						self       :  Optional[int])       ->   str:



						"""simple docstring"""






						import bitsandbytes as bnb

						from transformers import TaForConditionalGeneration

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

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

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







class       __lowerCAmelCase   (      lowerCamelCase_						):



			def    _lowerCamelCase	(						self       :  List[Any])       ->   List[str]:



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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   int:



						"""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       :  Union[str, Any])       ->   List[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       __lowerCAmelCase   (      lowerCamelCase_						):



			def    _lowerCamelCase	(						self       :  Tuple)       ->   Optional[Any]:



						"""simple docstring"""






						super().setUp()



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



						"""simple docstring"""






						del self.pipe

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



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



						"""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   (      lowerCamelCase_						):



			def    _lowerCamelCase	(						self       :  List[Any])       ->   str:



						"""simple docstring"""






						super().setUp()



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



						"""simple docstring"""






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







class       __lowerCAmelCase   (      lowerCamelCase_						):



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



						"""simple docstring"""






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



			def    _lowerCamelCase	(						self       :  Tuple)       ->   Optional[int]:



						"""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=_UpperCAmelCase)

						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(_UpperCAmelCase)):
												_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(**_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        =					"""gpt2-xl"""
			UpperCamelCase        =					3.3_191_854_854_152_187











from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


UpperCAmelCase__       =	{
    "configuration_clip": [
        "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "CLIPConfig",
        "CLIPOnnxConfig",
        "CLIPTextConfig",
        "CLIPVisionConfig",
    ],
    "processing_clip": ["CLIPProcessor"],
    "tokenization_clip": ["CLIPTokenizer"],
}

try:
							if not is_tokenizers_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	["CLIPTokenizerFast"]

try:
							if not is_vision_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	["CLIPFeatureExtractor"]
							UpperCAmelCase__       =	["CLIPImageProcessor"]

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	[
							    "CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "CLIPModel",
							    "CLIPPreTrainedModel",
							    "CLIPTextModel",
							    "CLIPTextModelWithProjection",
							    "CLIPVisionModel",
							    "CLIPVisionModelWithProjection",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	[
							    "TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "TFCLIPModel",
							    "TFCLIPPreTrainedModel",
							    "TFCLIPTextModel",
							    "TFCLIPVisionModel",
							]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase__       =	[
							    "FlaxCLIPModel",
							    "FlaxCLIPPreTrainedModel",
							    "FlaxCLIPTextModel",
							    "FlaxCLIPTextPreTrainedModel",
							    "FlaxCLIPVisionModel",
							    "FlaxCLIPVisionPreTrainedModel",
							]


if TYPE_CHECKING:
							from .configuration_clip import (
							    CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    CLIPConfig,
							    CLIPOnnxConfig,
							    CLIPTextConfig,
							    CLIPVisionConfig,
							)
							from .processing_clip import CLIPProcessor
							from .tokenization_clip import CLIPTokenizer

							try:
														if not is_tokenizers_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .tokenization_clip_fast import CLIPTokenizerFast

							try:
														if not is_vision_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .feature_extraction_clip import CLIPFeatureExtractor
														from .image_processing_clip import CLIPImageProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_clip import (
														    CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
														    CLIPModel,
														    CLIPPreTrainedModel,
														    CLIPTextModel,
														    CLIPTextModelWithProjection,
														    CLIPVisionModel,
														    CLIPVisionModelWithProjection,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_clip import (
														    TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
														    TFCLIPModel,
														    TFCLIPPreTrainedModel,
														    TFCLIPTextModel,
														    TFCLIPVisionModel,
														)

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_clip import (
														    FlaxCLIPModel,
														    FlaxCLIPPreTrainedModel,
														    FlaxCLIPTextModel,
														    FlaxCLIPTextPreTrainedModel,
														    FlaxCLIPVisionModel,
														    FlaxCLIPVisionPreTrainedModel,
														)


else:
							import sys

							UpperCAmelCase__       =	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
