

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







							'''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      ):




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




									def    a							( self				,       snake_case				,       *snake_case				,       **snake_case	):
																return self.module(snake_case				,       *snake_case				,       **snake_case	) + self.adapter(snake_case	)







@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			lowercase    (	unittest.TestCase      ):
				# 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 :			Optional[Any]	      =      '''bigscience/bloom-1b7'''

				# Constant values
				__SCREAMING_SNAKE_CASE :			Dict	      =      2.109659552692574

				__SCREAMING_SNAKE_CASE :			List[Any]	      =      '''Hello my name is'''
				__SCREAMING_SNAKE_CASE :			Tuple	      =      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 :			Optional[Any]	      =      10




				def    a							( self	):
											# Models and tokenizer
											snake_case_          =   AutoTokenizer.from_pretrained(self.model_name	)







class 			lowercase    (	lowercase_      ):




				def    a							( self	):
											super().setUp()

											# Models and tokenizer
											snake_case_          =   AutoModelForCausalLM.from_pretrained(
											    self.model_name				,       torch_dtype=torch.floataa				,       device_map='auto'	)
											snake_case_          =   AutoModelForCausalLM.from_pretrained(self.model_name				,       load_in_abit=snake_case				,       device_map='auto'	)




				def    a							( self	):
											del self.model_fpaa
											del self.model_abit

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




				def    a							( self	):
											snake_case_          =   self.model_abit.config

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

											snake_case_          =   config.to_dict()
											snake_case_          =   config.to_diff_dict()

											snake_case_          =   config.to_json_string()




				def    a							( self	):
											from bitsandbytes.nn import Paramsabit

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

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




				def    a							( self	):
											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(snake_case				,       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    a							( self	):
											snake_case_          =   self.tokenizer(self.input_text				,       return_tensors='pt'	)
											snake_case_          =   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=snake_case	)				,       self.EXPECTED_OUTPUTS	)




				def    a							( self	):
											snake_case_          =   BitsAndBytesConfig()
											snake_case_          =   True

											snake_case_          =   AutoModelForCausalLM.from_pretrained(
											    self.model_name				,       quantization_config=snake_case				,       device_map='auto'	)

											snake_case_          =   self.tokenizer(self.input_text				,       return_tensors='pt'	)
											snake_case_          =   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=snake_case	)				,       self.EXPECTED_OUTPUTS	)




				def    a							( self	):
											with self.assertRaises(snake_case	), tempfile.TemporaryDirectory() as tmpdirname:
																		self.model_abit.save_pretrained(snake_case	)




				def    a							( self	):
											snake_case_          =   BitsAndBytesConfig()

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




				def    a							( self	):
											with self.assertRaises(snake_case	):
																		# Tries with `str`
																		self.model_abit.to('cpu'	)

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

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

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

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

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

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

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

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

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




				def    a							( self	):
											snake_case_          =   AutoModelForSeqaSeqLM.from_pretrained('t5-small'				,       load_in_abit=snake_case				,       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      ):




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




				def    a							( self	):
											gc.collect()
											torch.cuda.empty_cache()




				def    a							( self	):
											from transformers import TaForConditionalGeneration

											snake_case_          =   TaForConditionalGeneration._keep_in_fpaa_modules
											snake_case_          =   None

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

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




				def    a							( self	):
											import bitsandbytes as bnb

											from transformers import TaForConditionalGeneration

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

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

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







class 			lowercase    (	lowercase_      ):




				def    a							( self	):
											super().setUp()
											# model_name
											snake_case_          =   'bigscience/bloom-560m'
											snake_case_          =   't5-small'

											# Different types of model

											snake_case_          =   AutoModel.from_pretrained(self.model_name				,       load_in_abit=snake_case				,       device_map='auto'	)
											# Sequence classification model
											snake_case_          =   AutoModelForSequenceClassification.from_pretrained(
											    self.model_name				,       load_in_abit=snake_case				,       device_map='auto'	)
											# CausalLM model
											snake_case_          =   AutoModelForCausalLM.from_pretrained(self.model_name				,       load_in_abit=snake_case				,       device_map='auto'	)
											# Seq2seq model
											snake_case_          =   AutoModelForSeqaSeqLM.from_pretrained(
											    self.seq_to_seq_name				,       load_in_abit=snake_case				,       device_map='auto'	)




				def    a							( self	):
											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    a							( self	):
											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    (	lowercase_      ):




				def    a							( self	):
											super().setUp()




				def    a							( self	):
											del self.pipe

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




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







@require_torch_multi_gpu
class 			lowercase    (	lowercase_      ):




				def    a							( self	):
											super().setUp()




				def    a							( self	):
											snake_case_          =   AutoModelForCausalLM.from_pretrained(
											    self.model_name				,       load_in_abit=snake_case				,       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
											snake_case_          =   self.tokenizer(self.input_text				,       return_tensors='pt'	)

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







class 			lowercase    (	lowercase_      ):




				def    a							( self	):
											snake_case_          =   'facebook/opt-350m'
											super().setUp()




				def    a							( self	):
											if version.parse(importlib.metadata.version('bitsandbytes'	)	) < version.parse('0.37.0'	):
																		return

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

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

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

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

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

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

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







class 			lowercase    (	lowercase_      ):
				__SCREAMING_SNAKE_CASE :			Optional[Any]	      =      '''gpt2-xl'''
				__SCREAMING_SNAKE_CASE :			int	      =      3.3191854854152187







import os
import sys
import unittest


_UpperCAmelCase		:       str						    =	os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, """utils"""))

import get_test_info  # noqa: E402
from get_test_info import (  # noqa: E402
    get_model_to_test_mapping,
    get_model_to_tester_mapping,
    get_test_to_tester_mapping,
)


_UpperCAmelCase		:       List[Any]						    =	os.path.join("""tests""", """models""", """bert""", """test_modeling_bert.py""")
_UpperCAmelCase		:       Dict						    =	os.path.join("""tests""", """models""", """blip""", """test_modeling_blip.py""")







class 			lowercase    (	unittest.TestCase      ):




				def    a							( self	):
											snake_case_          =   get_test_to_tester_mapping(snake_case	)
											snake_case_          =   get_test_to_tester_mapping(snake_case	)

											snake_case_          =   {'BertModelTest': 'BertModelTester'}

											snake_case_          =   {
											    'BlipModelTest': 'BlipModelTester',
											    'BlipTextImageModelTest': 'BlipTextImageModelsModelTester',
											    'BlipTextModelTest': 'BlipTextModelTester',
											    'BlipTextRetrievalModelTest': 'BlipTextRetrievalModelTester',
											    'BlipVQAModelTest': 'BlipVQAModelTester',
											    'BlipVisionModelTest': 'BlipVisionModelTester',
											}

											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)
											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)




				def    a							( self	):
											snake_case_          =   get_model_to_test_mapping(snake_case	)
											snake_case_          =   get_model_to_test_mapping(snake_case	)

											snake_case_          =   {
											    'BertForMaskedLM': ['BertModelTest'],
											    'BertForMultipleChoice': ['BertModelTest'],
											    'BertForNextSentencePrediction': ['BertModelTest'],
											    'BertForPreTraining': ['BertModelTest'],
											    'BertForQuestionAnswering': ['BertModelTest'],
											    'BertForSequenceClassification': ['BertModelTest'],
											    'BertForTokenClassification': ['BertModelTest'],
											    'BertLMHeadModel': ['BertModelTest'],
											    'BertModel': ['BertModelTest'],
											}

											snake_case_          =   {
											    'BlipForConditionalGeneration': ['BlipTextImageModelTest'],
											    'BlipForImageTextRetrieval': ['BlipTextRetrievalModelTest'],
											    'BlipForQuestionAnswering': ['BlipVQAModelTest'],
											    'BlipModel': ['BlipModelTest'],
											    'BlipTextModel': ['BlipTextModelTest'],
											    'BlipVisionModel': ['BlipVisionModelTest'],
											}

											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)
											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)




				def    a							( self	):
											snake_case_          =   get_model_to_tester_mapping(snake_case	)
											snake_case_          =   get_model_to_tester_mapping(snake_case	)

											snake_case_          =   {
											    'BertForMaskedLM': ['BertModelTester'],
											    'BertForMultipleChoice': ['BertModelTester'],
											    'BertForNextSentencePrediction': ['BertModelTester'],
											    'BertForPreTraining': ['BertModelTester'],
											    'BertForQuestionAnswering': ['BertModelTester'],
											    'BertForSequenceClassification': ['BertModelTester'],
											    'BertForTokenClassification': ['BertModelTester'],
											    'BertLMHeadModel': ['BertModelTester'],
											    'BertModel': ['BertModelTester'],
											}

											snake_case_          =   {
											    'BlipForConditionalGeneration': ['BlipTextImageModelsModelTester'],
											    'BlipForImageTextRetrieval': ['BlipTextRetrievalModelTester'],
											    'BlipForQuestionAnswering': ['BlipVQAModelTester'],
											    'BlipModel': ['BlipModelTester'],
											    'BlipTextModel': ['BlipTextModelTester'],
											    'BlipVisionModel': ['BlipVisionModelTester'],
											}

											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)
											self.assertEqual(get_test_info.to_json(snake_case	)				,       snake_case	)




