

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	(      SCREAMING_SNAKE_CASE      ):
	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				,      lowerCamelCase__				,      lowerCamelCase__						):
											super().__init__()
											UpperCAmelCase__:	Any									= module
											UpperCAmelCase__:	Dict									= nn.Sequential(
											    nn.Linear(module.in_features				,      lowerCamelCase__				,      bias=lowerCamelCase__						)				,      nn.Linear(lowerCamelCase__				,      module.out_features				,      bias=lowerCamelCase__						)				,      )
											UpperCAmelCase__:	Tuple									= (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__						):
											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'''



					__magic_name__        =   "bigscience/bloom-1b7"

					# Constant values
					__magic_name__        =   2.109_6595_5269_2574

					__magic_name__        =   "Hello my name is"
					__magic_name__        =   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"		)
					__magic_name__        =   1_0





					def    _UpperCAmelCase    (				self						):
						# Models and tokenizer
						UpperCAmelCase__:	Any									= AutoTokenizer.from_pretrained(self.model_name						)

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



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

						# Models and tokenizer
						UpperCAmelCase__:	List[str]									= AutoModelForCausalLM.from_pretrained(
						    self.model_name				,      torch_dtype=torch.floataa				,      device_map="auto"						)
						UpperCAmelCase__:	str									= AutoModelForCausalLM.from_pretrained(self.model_name				,      load_in_abit=lowerCamelCase__				,      device_map="auto"						)





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

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





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	int									= self.model_abit.config

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

						UpperCAmelCase__:	Any									= config.to_dict()
						UpperCAmelCase__:	Tuple									= config.to_diff_dict()

						UpperCAmelCase__:	int									= config.to_json_string()





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

						UpperCAmelCase__:	int									= self.model_fpaa.get_memory_footprint()
						UpperCAmelCase__:	Any									= self.model_abit.get_memory_footprint()

						self.assertAlmostEqual(mem_fpaa / mem_abit				,      self.EXPECTED_RELATIVE_DIFFERENCE						)
						UpperCAmelCase__:	Optional[int]									= get_some_linear_layer(self.model_abit						)
						self.assertTrue(linear.weight.__class__ == Paramsabit						)





					def    _UpperCAmelCase    (				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(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						):
						UpperCAmelCase__:	str									= self.tokenizer(self.input_text				,      return_tensors="pt"						)
						UpperCAmelCase__:	Tuple									= self.model_abit.generate(input_ids=encoded_input["input_ids"].to(0						)				,      max_new_tokens=1_0						)

						self.assertIn(self.tokenizer.decode(output_sequences[0]				,      skip_special_tokens=lowerCamelCase__						)				,      self.EXPECTED_OUTPUTS						)





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Optional[Any]									= BitsAndBytesConfig()
						UpperCAmelCase__:	Dict									= True

						UpperCAmelCase__:	str									= AutoModelForCausalLM.from_pretrained(
						    self.model_name				,      quantization_config=lowerCamelCase__				,      device_map="auto"						)

						UpperCAmelCase__:	Any									= self.tokenizer(self.input_text				,      return_tensors="pt"						)
						UpperCAmelCase__:	Optional[int]									= model_abit_from_config.generate(
						    input_ids=encoded_input["input_ids"].to(0						)				,      max_new_tokens=1_0						)

						self.assertIn(self.tokenizer.decode(output_sequences[0]				,      skip_special_tokens=lowerCamelCase__						)				,      self.EXPECTED_OUTPUTS						)





					def    _UpperCAmelCase    (				self						):
						with self.assertRaises(lowerCamelCase__						), tempfile.TemporaryDirectory() as tmpdirname:
							self.model_abit.save_pretrained(lowerCamelCase__						)





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	int									= BitsAndBytesConfig()

						with self.assertRaises(lowerCamelCase__						):
							UpperCAmelCase__:	Dict									= AutoModelForCausalLM.from_pretrained(
							    self.model_name				,      quantization_config=lowerCamelCase__				,      load_in_abit=lowerCamelCase__				,      device_map="auto"				,      bnb_abit_quant_type="nf4"				,      )





					def    _UpperCAmelCase    (				self						):
						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
						UpperCAmelCase__:	Optional[Any]									= self.tokenizer(self.input_text				,      return_tensors="pt"						)

						UpperCAmelCase__:	Dict									= self.model_fpaa.to(torch.floataa						)
						UpperCAmelCase__:	Any									= self.model_fpaa.generate(input_ids=encoded_input["input_ids"].to(0						)				,      max_new_tokens=1_0						)

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

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

						# Check this does not throw an error
						UpperCAmelCase__:	Union[str, Any]									= self.model_fpaa.float()





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Union[str, Any]									= 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    _UpperCAmelCase    (				cls						):
						UpperCAmelCase__:	str									= "t5-small"
						UpperCAmelCase__:	int									= "google/flan-t5-small"  # flan-t5 uses dense-act instead of dense-relu-dense
						UpperCAmelCase__:	Optional[int]									= AutoTokenizer.from_pretrained(cls.model_name						)
						UpperCAmelCase__:	Dict									= "Translate in German: Hello, my dog is cute"





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





					def    _UpperCAmelCase    (				self						):
						from transformers import TaForConditionalGeneration

						UpperCAmelCase__:	Optional[int]									= TaForConditionalGeneration._keep_in_fpaa_modules
						UpperCAmelCase__:	List[str]									= None

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

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





					def    _UpperCAmelCase    (				self						):
						import bitsandbytes as bnb

						from transformers import TaForConditionalGeneration

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

						UpperCAmelCase__:	List[str]									= self.tokenizer(self.input_text				,      return_tensors="pt"						).to(0						)
						UpperCAmelCase__:	Dict									= model.generate(**lowerCamelCase__						)

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

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					def    _UpperCAmelCase    (				self						):
						super().setUp()
						# model_name
						UpperCAmelCase__:	int									= "bigscience/bloom-560m"
						UpperCAmelCase__:	List[Any]									= "t5-small"

						# Different types of model

						UpperCAmelCase__:	Any									= AutoModel.from_pretrained(self.model_name				,      load_in_abit=lowerCamelCase__				,      device_map="auto"						)
						# Sequence classification model
						UpperCAmelCase__:	str									= AutoModelForSequenceClassification.from_pretrained(
						    self.model_name				,      load_in_abit=lowerCamelCase__				,      device_map="auto"						)
						# CausalLM model
						UpperCAmelCase__:	Optional[Any]									= AutoModelForCausalLM.from_pretrained(self.model_name				,      load_in_abit=lowerCamelCase__				,      device_map="auto"						)
						# Seq2seq model
						UpperCAmelCase__:	List[str]									= AutoModelForSeqaSeqLM.from_pretrained(
						    self.seq_to_seq_name				,      load_in_abit=lowerCamelCase__				,      device_map="auto"						)





					def    _UpperCAmelCase    (				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    _UpperCAmelCase    (				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 __UpperCamelCase     (				_a		):






					'''simple docstring'''



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





					def    _UpperCAmelCase    (				self						):
						del self.pipe

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





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	List[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
						UpperCAmelCase__:	str									= 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						):
						super().setUp()





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

						# Second real batch
						UpperCAmelCase__:	Any									= model_parallel.generate(input_ids=encoded_input["input_ids"].to(0						)				,      max_new_tokens=1_0						)
						self.assertIn(self.tokenizer.decode(output_parallel[0]				,      skip_special_tokens=lowerCamelCase__						)				,      self.EXPECTED_OUTPUTS						)

class __UpperCamelCase     (				_a		):






					'''simple docstring'''



					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Optional[int]									= "facebook/opt-350m"
						super().setUp()





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

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

        # Step 2: add adapters
						for _, module in model.named_modules():
							if "OPTAttention" in repr(type(lowerCamelCase__						)						):
								UpperCAmelCase__:	Any									= LoRALayer(module.q_proj				,      rank=1_6						)
								UpperCAmelCase__:	List[str]									= LoRALayer(module.k_proj				,      rank=1_6						)
								UpperCAmelCase__:	List[Any]									= LoRALayer(module.v_proj				,      rank=1_6						)

        # Step 3: dummy batch
						UpperCAmelCase__:	Optional[int]									= self.tokenizer("Test batch "				,      return_tensors="pt"						).to(0						)

						# Step 4: Check if the gradient is not None
						with torch.cuda.amp.autocast():
							UpperCAmelCase__:	List[str]									= 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     (				_a		):






					'''simple docstring'''



					__magic_name__        =   "gpt2-xl"
					__magic_name__        =   3.3191_8548_5415_2187


_lowerCAmelCase  :      int			 ="""
# Installazione di Transformers
! pip install transformers datasets
# Per installare dalla fonte invece dell'ultima versione rilasciata, commenta il comando sopra e
# rimuovi la modalità commento al comando seguente.
# ! pip install git+https://github.com/huggingface/transformers.git
"""

_lowerCAmelCase  :      List[str]			 =[{"""type""": """code""", """content""": INSTALL_CONTENT}]
_lowerCAmelCase  :      int			 ={
    """{processor_class}""": """FakeProcessorClass""",
    """{model_class}""": """FakeModelClass""",
    """{object_class}""": """FakeObjectClass""",
}