





'''simple docstring'''




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


		"""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      UpperCamelCase__					(					nn.Module			):




												"""simple docstring"""


												def __init__(						self				:			Optional[Any]      ,						lowerCamelCase_				:			nn.Module      ,						lowerCamelCase_				:			int						):

														'''simple docstring'''



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

														'''simple docstring'''



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




@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class      UpperCamelCase__					(					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
					SCREAMING_SNAKE_CASE__				=	'''bigscience/bloom-1b7'''

					# Constant values
					SCREAMING_SNAKE_CASE__				=	2.109659552692574

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


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



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







class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			Any						):

							'''simple docstring'''



							super().setUp()

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


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							del self.model_fpaa
							del self.model_abit

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


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			self.model_abit.config

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

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

							SCREAMING_SNAKE_CASE :				int		=			config.to_json_string()


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							from bitsandbytes.nn import Paramsabit

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

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


					def lowerCamelCase_			(						self				:			Optional[Any]						):

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

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				List[Any]		=			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=lowerCamelCase_						)      ,						self.EXPECTED_OUTPUTS						)


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			BitsAndBytesConfig()
							SCREAMING_SNAKE_CASE :				Optional[int]		=			True

							SCREAMING_SNAKE_CASE :				List[Any]		=			AutoModelForCausalLM.from_pretrained(
							    self.model_name      ,						quantization_config=lowerCamelCase_      ,						device_map="""auto"""						)

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

							'''simple docstring'''



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


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				str		=			BitsAndBytesConfig()

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


					def lowerCamelCase_			(						self				:			Optional[int]						):

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

							SCREAMING_SNAKE_CASE :				List[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 :				str		=			self.model_fpaa.to("""cpu"""						)

							# Check this does not throw an error
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.model_fpaa.half()

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



					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			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      UpperCamelCase__					(					unittest.TestCase			):




					"""simple docstring"""


					@classmethod
					def lowerCamelCase_			(						cls				:			str						):

							'''simple docstring'''



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


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



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


					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							from transformers import TaForConditionalGeneration

							SCREAMING_SNAKE_CASE :				List[str]		=			TaForConditionalGeneration._keep_in_fpaa_modules
							SCREAMING_SNAKE_CASE :				int		=			None

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

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



					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							import bitsandbytes as bnb

							from transformers import TaForConditionalGeneration

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

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

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







class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



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

							# Different types of model

							SCREAMING_SNAKE_CASE :				Dict		=			AutoModel.from_pretrained(self.model_name      ,						load_in_abit=lowerCamelCase_      ,						device_map="""auto"""						)
							# Sequence classification model
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			AutoModelForSequenceClassification.from_pretrained(
							    self.model_name      ,						load_in_abit=lowerCamelCase_      ,						device_map="""auto"""						)
							# CausalLM model
							SCREAMING_SNAKE_CASE :				Tuple		=			AutoModelForCausalLM.from_pretrained(self.model_name      ,						load_in_abit=lowerCamelCase_      ,						device_map="""auto"""						)
							# Seq2seq model
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			AutoModelForSeqaSeqLM.from_pretrained(
							    self.seq_to_seq_name      ,						load_in_abit=lowerCamelCase_      ,						device_map="""auto"""						)


					def lowerCamelCase_			(						self				:			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				:			int						):

							'''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__					(					lowercase_			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			Dict						):

							'''simple docstring'''



							super().setUp()


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							del self.pipe

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



					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



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







@require_torch_multi_gpu
class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							super().setUp()



					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



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

							# Second real batch
							SCREAMING_SNAKE_CASE :				Dict		=			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      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



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



					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



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

							# Step 1: freeze all parameters
							SCREAMING_SNAKE_CASE :				List[str]		=			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():
									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(lowerCamelCase_						)						):
											SCREAMING_SNAKE_CASE :				int		=			LoRALayer(module.q_proj      ,						rank=16						)
											SCREAMING_SNAKE_CASE :				str		=			LoRALayer(module.k_proj      ,						rank=16						)
											SCREAMING_SNAKE_CASE :				List[Any]		=			LoRALayer(module.v_proj      ,						rank=16						)

        # Step 3: dummy batch
							SCREAMING_SNAKE_CASE :				List[str]		=			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 :				Dict		=			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      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	'''gpt2-xl'''
					SCREAMING_SNAKE_CASE__				=	3.3191854854152187










'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCAmelCase            =       logging.get_logger(__name__)

__UpperCAmelCase            =       {
    """google/fnet-base""": """https://huggingface.co/google/fnet-base/resolve/main/config.json""",
    """google/fnet-large""": """https://huggingface.co/google/fnet-large/resolve/main/config.json"""
    # See all FNet models at https://huggingface.co/models?filter=fnet
}




class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""
					SCREAMING_SNAKE_CASE__				=	'''fnet'''


					def __init__(						self				:			Any      ,						lowerCamelCase_				:			List[str]=3_20_00      ,						lowerCamelCase_				:			List[Any]=7_68      ,						lowerCamelCase_				:			Union[str, Any]=12      ,						lowerCamelCase_				:			Optional[Any]=30_72      ,						lowerCamelCase_				:			Optional[Any]="gelu_new"      ,						lowerCamelCase_				:			Any=0.1      ,						lowerCamelCase_				:			Union[str, Any]=5_12      ,						lowerCamelCase_				:			Optional[int]=4      ,						lowerCamelCase_				:			Optional[int]=0.02      ,						lowerCamelCase_				:			str=1e-12      ,						lowerCamelCase_				:			Any=False      ,						lowerCamelCase_				:			str=5_12      ,						lowerCamelCase_				:			str=3      ,						lowerCamelCase_				:			int=1      ,						lowerCamelCase_				:			Optional[Any]=2      ,						**lowerCamelCase_				:			Optional[int]      ,						):

							'''simple docstring'''



							super().__init__(pad_token_id=lowerCamelCase_      ,						bos_token_id=lowerCamelCase_      ,						eos_token_id=lowerCamelCase_      ,						**lowerCamelCase_						)

							SCREAMING_SNAKE_CASE :				Dict		=			vocab_size
							SCREAMING_SNAKE_CASE :				Dict		=			max_position_embeddings
							SCREAMING_SNAKE_CASE :				str		=			hidden_size
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			num_hidden_layers
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			intermediate_size
							SCREAMING_SNAKE_CASE :				int		=			hidden_act
							SCREAMING_SNAKE_CASE :				Dict		=			hidden_dropout_prob
							SCREAMING_SNAKE_CASE :				Tuple		=			initializer_range
							SCREAMING_SNAKE_CASE :				List[Any]		=			type_vocab_size
							SCREAMING_SNAKE_CASE :				Any		=			layer_norm_eps
							SCREAMING_SNAKE_CASE :				int		=			use_tpu_fourier_optimizations
							SCREAMING_SNAKE_CASE :				List[Any]		=			tpu_short_seq_length



