




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


__A        =    {"configuration_speech_encoder_decoder": ["SpeechEncoderDecoderConfig"]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__A        =    ["SpeechEncoderDecoderModel"]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__A        =    ["FlaxSpeechEncoderDecoderModel"]

if TYPE_CHECKING:
					from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
					import sys

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







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(      A       :    List[str]       )		->       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 	_lowercase							(   nn.Module	):





					"""simple docstring"""




					def __init__(self		,       lowerCamelCase_		,       lowerCamelCase_		):
										"""simple docstring"""

										super().__init__()
										a					     =   module
										a					     =   nn.Sequential(
										    nn.Linear(module.in_features		,       lowerCamelCase_		,       bias=lowerCamelCase_		)		,       nn.Linear(lowerCamelCase_		,       module.out_features		,       bias=lowerCamelCase_		)		,       )
										a					     =   (2.0 / (5 * min(module.in_features		,       module.out_features		))) ** 0.5
										nn.init.normal_(self.adapter[0].weight		,       std=lowerCamelCase_		)
										nn.init.zeros_(self.adapter[1].weight		)
										self.adapter.to(module.weight.device		)




					def      UpperCamelCase_       (self		,       lowerCamelCase_		,       *lowerCamelCase_		,       **lowerCamelCase_		):
										"""simple docstring"""

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



@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 	_lowercase							(   unittest.TestCase	):





	"""simple docstring"""
	# 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
	__A							=							"bigscience/bloom-1b7"

	# Constant values
	__A							=							2.109_659_552_692_574

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

						a					     =   AutoTokenizer.from_pretrained(self.model_name		)




class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						super().setUp()

						# Models and tokenizer
						a					     =   AutoModelForCausalLM.from_pretrained(
						    self.model_name		,       torch_dtype=torch.floataa		,       device_map="auto"		)
						a					     =   AutoModelForCausalLM.from_pretrained(self.model_name		,       load_in_abit=lowerCamelCase_		,       device_map="auto"		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						del self.model_fpaa
						del self.model_abit

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   self.model_abit.config

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

						a					     =   config.to_dict()
						a					     =   config.to_diff_dict()

						a					     =   config.to_json_string()




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						from bitsandbytes.nn import Paramsabit

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

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




	def      UpperCamelCase_       (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(lowerCamelCase_		,       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		):
						"""simple docstring"""

						a					     =   self.tokenizer(self.input_text		,       return_tensors="pt"		)
						a					     =   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=lowerCamelCase_		)		,       self.EXPECTED_OUTPUTS		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   BitsAndBytesConfig()
						a					     =   True

						a					     =   AutoModelForCausalLM.from_pretrained(
						    self.model_name		,       quantization_config=lowerCamelCase_		,       device_map="auto"		)

						a					     =   self.tokenizer(self.input_text		,       return_tensors="pt"		)
						a					     =   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=lowerCamelCase_		)		,       self.EXPECTED_OUTPUTS		)




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   BitsAndBytesConfig()

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						with self.assertRaises(lowerCamelCase_		):
											# Tries with `str`
											self.model_abit.to("cpu"		)

						with self.assertRaises(lowerCamelCase_		):
											# Tries with a `dtype``
											self.model_abit.to(torch.floataa		)

						with self.assertRaises(lowerCamelCase_		):
											# Tries with a `device`
											self.model_abit.to(torch.device("cuda:0"		)		)

						with self.assertRaises(lowerCamelCase_		):
											# Tries with a `device`
											self.model_abit.float()

						with self.assertRaises(lowerCamelCase_		):
											# Tries with a `device`
											self.model_abit.half()

						# Test if we did not break anything
						a					     =   self.tokenizer(self.input_text		,       return_tensors="pt"		)

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

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

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

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   AutoModelForSeqaSeqLM.from_pretrained("t5-small"		,       load_in_abit=lowerCamelCase_		,       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 	_lowercase							(   unittest.TestCase	):





	"""simple docstring"""




	@classmethod
	def      UpperCamelCase_       (cls		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						from transformers import TaForConditionalGeneration

						a					     =   TaForConditionalGeneration._keep_in_fpaa_modules
						a					     =   None

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

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						import bitsandbytes as bnb

						from transformers import TaForConditionalGeneration

						# test with `t5-small`
						a					     =   TaForConditionalGeneration.from_pretrained(self.model_name		,       load_in_abit=lowerCamelCase_		,       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		)		)

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

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




class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						super().setUp()
						# model_name
						a					     =   "bigscience/bloom-560m"
						a					     =   "t5-small"

						# Different types of model

						a					     =   AutoModel.from_pretrained(self.model_name		,       load_in_abit=lowerCamelCase_		,       device_map="auto"		)
						# Sequence classification model
						a					     =   AutoModelForSequenceClassification.from_pretrained(
						    self.model_name		,       load_in_abit=lowerCamelCase_		,       device_map="auto"		)
						# CausalLM model
						a					     =   AutoModelForCausalLM.from_pretrained(self.model_name		,       load_in_abit=lowerCamelCase_		,       device_map="auto"		)
						# Seq2seq model
						a					     =   AutoModelForSeqaSeqLM.from_pretrained(
						    self.seq_to_seq_name		,       load_in_abit=lowerCamelCase_		,       device_map="auto"		)




	def      UpperCamelCase_       (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      UpperCamelCase_       (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 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						super().setUp()




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						del self.pipe

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




@require_torch_multi_gpu
class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						super().setUp()




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

						a					     =   AutoModelForCausalLM.from_pretrained(
						    self.model_name		,       load_in_abit=lowerCamelCase_		,       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
						a					     =   self.tokenizer(self.input_text		,       return_tensors="pt"		)

						# Second real batch
						a					     =   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=lowerCamelCase_		)		,       self.EXPECTED_OUTPUTS		)




class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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




	def      UpperCamelCase_       (self		):
						"""simple docstring"""

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

						# Step 1: freeze all parameters
						a					     =   AutoModelForCausalLM.from_pretrained(self.model_name		,       load_in_abit=lowerCamelCase_		)

						self.assertEqual(set(model.hf_device_map.values()		)		,       {torch.cuda.current_device()}		)

						for param in model.parameters():
											a					     =   False  # freeze the model - train adapters later
											if param.ndim == 1:
																# cast the small parameters (e.g. layernorm) to fp32 for stability
																a					     =   param.data.to(torch.floataa		)

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

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

						# Step 4: Check if the gradient is not None
						with torch.cuda.amp.autocast():
											a					     =   model.forward(**lowerCamelCase_		)
											out.logits.norm().backward()

						for module in model.modules():
											if isinstance(lowerCamelCase_		,       lowerCamelCase_		):
																self.assertTrue(module.adapter[1].weight.grad is not None		)
																self.assertTrue(module.adapter[1].weight.grad.norm().item() > 0		)
											elif isinstance(lowerCamelCase_		,       nn.Embedding		):
																self.assertTrue(module.weight.grad is None		)




class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""
	__A							=							"gpt2-xl"
	__A							=							3.3_191_854_854_152_187
