

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):
	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				:			Union[str, Any]   ,   a				:			nn.Module   ,   a				:			int						)				->							Dict:

							"""simple docstring"""

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

							"""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'''







			lowerCamelCase__							     ='bigscience/bloom-1b7'

			# Constant values
			lowerCamelCase__							     =2.109659552692574

			lowerCamelCase__							     ='Hello my name is'
			lowerCamelCase__							     =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' )
			lowerCamelCase__							     =10
			def 				__UpperCamelCase			(      self				:			Dict						)				->							Optional[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							AutoTokenizer.from_pretrained(self.model_name						)






class _UpperCamelCase				(				__A ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			Dict						)				->							int:

				"""simple docstring"""

				super().setUp()

				# Models and tokenizer
				SCREAMING_SNAKE_CASE :   Dict					  =							AutoModelForCausalLM.from_pretrained(
				    self.model_name   ,   torch_dtype=torch.floataa   ,   device_map="auto"						)
				SCREAMING_SNAKE_CASE :   int					  =							AutoModelForCausalLM.from_pretrained(self.model_name   ,   load_in_abit=a   ,   device_map="auto"						)
			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Dict:

				"""simple docstring"""

				del self.model_fpaa
				del self.model_abit

				gc.collect()
				torch.cuda.empty_cache()
			def 				__UpperCamelCase			(      self				:			Tuple						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							self.model_abit.config

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

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							config.to_dict()
				SCREAMING_SNAKE_CASE :   str					  =							config.to_diff_dict()

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							config.to_json_string()
			def 				__UpperCamelCase			(      self				:			Dict						)				->							Optional[int]:

				"""simple docstring"""

				from bitsandbytes.nn import Paramsabit

				SCREAMING_SNAKE_CASE :   int					  =							self.model_fpaa.get_memory_footprint()
				SCREAMING_SNAKE_CASE :   int					  =							self.model_abit.get_memory_footprint()

				self.assertAlmostEqual(mem_fpaa / mem_abit   ,   self.EXPECTED_RELATIVE_DIFFERENCE						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							get_some_linear_layer(self.model_abit						)
				self.assertTrue(linear.weight.__class__ == Paramsabit						)
			def 				__UpperCamelCase			(      self				:			Optional[Any]						)				->							Optional[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 				__UpperCamelCase			(      self				:			Dict						)				->							Union[str, Any]:

				"""simple docstring"""

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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							BitsAndBytesConfig()
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							True

				SCREAMING_SNAKE_CASE :   int					  =							AutoModelForCausalLM.from_pretrained(
				    self.model_name   ,   quantization_config=a   ,   device_map="auto"						)

				SCREAMING_SNAKE_CASE :   int					  =							self.tokenizer(self.input_text   ,   return_tensors="pt"						)
				SCREAMING_SNAKE_CASE :   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=a						)   ,   self.EXPECTED_OUTPUTS						)
			def 				__UpperCamelCase			(      self				:			Any						)				->							str:

				"""simple docstring"""

				with self.assertRaises(a						), tempfile.TemporaryDirectory() as tmpdirname:
					self.model_abit.save_pretrained(a						)
			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							BitsAndBytesConfig()

				with self.assertRaises(a						):
					SCREAMING_SNAKE_CASE :   List[str]					  =							AutoModelForCausalLM.from_pretrained(
					    self.model_name   ,   quantization_config=a   ,   load_in_abit=a   ,   device_map="auto"   ,   bnb_abit_quant_type="nf4"   ,   )
			def 				__UpperCamelCase			(      self				:			Any						)				->							Tuple:

				"""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
				SCREAMING_SNAKE_CASE :   List[str]					  =							self.tokenizer(self.input_text   ,   return_tensors="pt"						)

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.model_fpaa.to(torch.floataa						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							self.model_fpaa.generate(input_ids=encoded_input["input_ids"].to(0						)   ,   max_new_tokens=10						)

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

				# Check this does not throw an error
				SCREAMING_SNAKE_CASE :   Dict					  =							self.model_fpaa.half()

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






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   int					  =							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 				__UpperCamelCase			(      cls				:			str						)				->							str:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[str]					  =							"t5-small"
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							"google/flan-t5-small"  # flan-t5 uses dense-act instead of dense-relu-dense
				SCREAMING_SNAKE_CASE :   Any					  =							AutoTokenizer.from_pretrained(cls.model_name						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							"Translate in German: Hello, my dog is cute"
			def 				__UpperCamelCase			(      self				:			Union[str, Any]						)				->							Any:

				"""simple docstring"""

				gc.collect()
				torch.cuda.empty_cache()
			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Union[str, Any]:

				"""simple docstring"""

				from transformers import TaForConditionalGeneration

				SCREAMING_SNAKE_CASE :   Any					  =							TaForConditionalGeneration._keep_in_fpaa_modules
				SCREAMING_SNAKE_CASE :   Dict					  =							None

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

				# test with `flan-t5-small`
				SCREAMING_SNAKE_CASE :   Dict					  =							TaForConditionalGeneration.from_pretrained(
				    self.dense_act_model_name   ,   load_in_abit=a   ,   device_map="auto"						)
				SCREAMING_SNAKE_CASE :   int					  =							self.tokenizer(self.input_text   ,   return_tensors="pt"						).to(0						)
				SCREAMING_SNAKE_CASE :   Optional[int]					  =							model.generate(**a						)
				SCREAMING_SNAKE_CASE :   Tuple					  =							modules






			def 				__UpperCamelCase			(      self				:			int						)				->							Dict:

				"""simple docstring"""

				import bitsandbytes as bnb

				from transformers import TaForConditionalGeneration

				# test with `t5-small`
				SCREAMING_SNAKE_CASE :   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						)						)

				SCREAMING_SNAKE_CASE :   List[Any]					  =							self.tokenizer(self.input_text   ,   return_tensors="pt"						).to(0						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							model.generate(**a						)

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






class _UpperCamelCase				(				__A ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			Any						)				->							Tuple:

				"""simple docstring"""

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

				# Different types of model

				SCREAMING_SNAKE_CASE :   str					  =							AutoModel.from_pretrained(self.model_name   ,   load_in_abit=a   ,   device_map="auto"						)
				# Sequence classification model
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							AutoModelForSequenceClassification.from_pretrained(
				    self.model_name   ,   load_in_abit=a   ,   device_map="auto"						)
				# CausalLM model
				SCREAMING_SNAKE_CASE :   List[Any]					  =							AutoModelForCausalLM.from_pretrained(self.model_name   ,   load_in_abit=a   ,   device_map="auto"						)
				# Seq2seq model
				SCREAMING_SNAKE_CASE :   List[str]					  =							AutoModelForSeqaSeqLM.from_pretrained(
				    self.seq_to_seq_name   ,   load_in_abit=a   ,   device_map="auto"						)
			def 				__UpperCamelCase			(      self				:			Optional[int]						)				->							Dict:

				"""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 				__UpperCamelCase			(      self				:			List[Any]						)				->							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 _UpperCamelCase				(				__A ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			Dict						)				->							Dict:

				"""simple docstring"""

				super().setUp()
			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							Optional[int]:

				"""simple docstring"""

				del self.pipe

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






			def 				__UpperCamelCase			(      self				:			Dict						)				->							Dict:

				"""simple docstring"""

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






@require_torch_multi_gpu
class _UpperCamelCase				(				__A ):
			'''simple docstring'''







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

				"""simple docstring"""

				super().setUp()






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[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
				SCREAMING_SNAKE_CASE :   Tuple					  =							self.tokenizer(self.input_text   ,   return_tensors="pt"						)

				# Second real batch
				SCREAMING_SNAKE_CASE :   Tuple					  =							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				(				__A ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			List[str]						)				->							Dict:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   str					  =							"facebook/opt-350m"
				super().setUp()






			def 				__UpperCamelCase			(      self				:			int						)				->							List[Any]:

				"""simple docstring"""

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

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

        # Step 2: add adapters
				for _, module in model.named_modules():
					if "OPTAttention" in repr(type(a						)						):
						SCREAMING_SNAKE_CASE :   Optional[Any]					  =							LoRALayer(module.q_proj   ,   rank=16						)
						SCREAMING_SNAKE_CASE :   Optional[int]					  =							LoRALayer(module.k_proj   ,   rank=16						)
						SCREAMING_SNAKE_CASE :   Optional[int]					  =							LoRALayer(module.v_proj   ,   rank=16						)

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

				# Step 4: Check if the gradient is not None
				with torch.cuda.amp.autocast():
					SCREAMING_SNAKE_CASE :   str					  =							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				(				__A ):
			'''simple docstring'''







			lowerCamelCase__							     ='gpt2-xl'
			lowerCamelCase__							     =3.3191854854152187


def 			lowerCamelCase__			(       _a				,      _a):
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							0
	while b > 0:
		if b & 1:
			res += a

		a += a
		b >>= 1

	return res

def 			lowerCamelCase__			(       _a				,      _a				,      _a):
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							0
	while b > 0:
		if b & 1:
			SCREAMING_SNAKE_CASE :   Optional[Any]					  =							((res % c) + (a % c)) % c

		a += a
		b >>= 1

	return res