






"""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  __lowerCamelCase					(					lowerCAmelCase__				):
				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       snake_case_       (    nn.Module       ):



						"""simple docstring"""

						def __init__(      self  ,					__a  ,					__a						):


										"""simple docstring"""

										super().__init__()
										A__							 =       module
										A__							 =       nn.Sequential(
										    nn.Linear(module.in_features  ,					__a  ,					bias=__a						)  ,					nn.Linear(__a  ,					module.out_features  ,					bias=__a						)  ,					)
										A__							 =       (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  ,					__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       snake_case_       (    unittest.TestCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							Any   			=					"""bigscience/bloom-1b7"""

	# Constant values
	SCREAMING_SNAKE_CASE_:							List[Any]   			=					2.109659552692574

	SCREAMING_SNAKE_CASE_:							Optional[int]   			=					"""Hello my name is"""
	SCREAMING_SNAKE_CASE_:							Union[str, Any]   			=					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_:							List[Any]   			=					10


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       AutoTokenizer.from_pretrained(self.model_name						)





class       snake_case_       (    _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=__a  ,					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(__a  ,					'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(__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						):


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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       BitsAndBytesConfig()
					A__							 =       True

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

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


	def       _UpperCAmelCase						(      self						):


					"""simple docstring"""

					A__							 =       BitsAndBytesConfig()

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


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

					# test with `flan-t5-small`
					A__							 =       TaForConditionalGeneration.from_pretrained(
					    self.dense_act_model_name  ,					load_in_abit=__a  ,					device_map='auto'						)
					A__							 =       self.tokenizer(self.input_text  ,					return_tensors='pt'						).to(0						)
					A__							 =       model.generate(**__a						)
					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=__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						)						)

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

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





class       snake_case_       (    _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=__a  ,					device_map='auto'						)
					# Sequence classification model
					A__							 =       AutoModelForSequenceClassification.from_pretrained(
					    self.model_name  ,					load_in_abit=__a  ,					device_map='auto'						)
					# CausalLM model
					A__							 =       AutoModelForCausalLM.from_pretrained(self.model_name  ,					load_in_abit=__a  ,					device_map='auto'						)
					# Seq2seq model
					A__							 =       AutoModelForSeqaSeqLM.from_pretrained(
					    self.seq_to_seq_name  ,					load_in_abit=__a  ,					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       snake_case_       (    _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       snake_case_       (    _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=__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
					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=__a						)  ,					self.EXPECTED_OUTPUTS						)





class       snake_case_       (    _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=__a						)

					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(__a						)						):
													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(**__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       snake_case_       (    _lowerCamelCase       ):



	"""simple docstring"""

	SCREAMING_SNAKE_CASE_:							Tuple   			=					"""gpt2-xl"""
	SCREAMING_SNAKE_CASE_:							Dict   			=					3.3191854854152187










"""simple docstring"""






from maths.prime_check import is_prime



def  __lowerCamelCase					(					lowerCAmelCase__				):
				if not isinstance(lowerCAmelCase__       ,lowerCAmelCase__				):
								A__							 =       f'''Input value of [number={number}] must be an integer'''
								raise TypeError(lowerCAmelCase__				)
				if is_prime(lowerCAmelCase__				) and is_prime(number + 2				):
								return number + 2
				else:
								return -1


if __name__ == "__main__":
					import doctest

					doctest.testmod()


