



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 snake_case__  (			__SCREAMING_SNAKE_CASE )		->							Any:
		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 	lowerCamelCase      (						nn.Module  ):







								'''simple docstring'''







								def __init__(			self    ,	lowerCAmelCase    ,	lowerCAmelCase     ):
										super().__init__()
										UpperCAmelCase_						=				module
										UpperCAmelCase_						=				nn.Sequential(
										    nn.Linear(module.in_features    ,	lowerCAmelCase    ,	bias=lowerCAmelCase     )    ,	nn.Linear(lowerCAmelCase    ,	module.out_features    ,	bias=lowerCAmelCase     )    ,	)
										UpperCAmelCase_						=				(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    A__		(			self    ,	lowerCAmelCase    ,	*lowerCAmelCase    ,	**lowerCAmelCase     ):
										return self.module(lowerCAmelCase    ,	*lowerCAmelCase    ,	**lowerCAmelCase     ) + self.adapter(lowerCAmelCase     )



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







					'''simple docstring'''







					lowerCAmelCase_				:    Any									=  'bigscience/bloom-1b7'

					# Constant values
					lowerCAmelCase_				:    Union[str, Any]									=  2.1_0_9_6_5_9_5_5_2_6_9_2_5_7_4

					lowerCAmelCase_				:    List[Any]									=  'Hello my name is'
					lowerCAmelCase_				:    List[str]									=  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'  )
					lowerCAmelCase_				:    Any									=  10



					def    A__		(			self     ):
							# Models and tokenizer
							UpperCAmelCase_						=				AutoTokenizer.from_pretrained(self.model_name     )





class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







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

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



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

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



					def    A__		(			self     ):
							UpperCAmelCase_						=				self.model_abit.config

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

							UpperCAmelCase_						=				config.to_dict()
							UpperCAmelCase_						=				config.to_diff_dict()

							UpperCAmelCase_						=				config.to_json_string()



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

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

							self.assertAlmostEqual(mem_fpaa / mem_abit    ,	self.EXPECTED_RELATIVE_DIFFERENCE     )
							UpperCAmelCase_						=				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(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    A__		(			self     ):
							UpperCAmelCase_						=				self.tokenizer(self.input_text    ,	return_tensors="pt"     )
							UpperCAmelCase_						=				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    A__		(			self     ):
							UpperCAmelCase_						=				BitsAndBytesConfig()
							UpperCAmelCase_						=				True

							UpperCAmelCase_						=				AutoModelForCausalLM.from_pretrained(
							    self.model_name    ,	quantization_config=lowerCAmelCase    ,	device_map="auto"     )

							UpperCAmelCase_						=				self.tokenizer(self.input_text    ,	return_tensors="pt"     )
							UpperCAmelCase_						=				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    A__		(			self     ):
							with self.assertRaises(lowerCAmelCase     ), tempfile.TemporaryDirectory() as tmpdirname:
									self.model_abit.save_pretrained(lowerCAmelCase     )



					def    A__		(			self     ):
							UpperCAmelCase_						=				BitsAndBytesConfig()

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



					def    A__		(			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_						=				self.tokenizer(self.input_text    ,	return_tensors="pt"     )

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

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

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

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



					def    A__		(			self     ):
							UpperCAmelCase_						=				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 	lowerCamelCase      (						unittest.TestCase  ):







					'''simple docstring'''







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



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



					def    A__		(			self     ):
							from transformers import TaForConditionalGeneration

							UpperCAmelCase_						=				TaForConditionalGeneration._keep_in_fpaa_modules
							UpperCAmelCase_						=				None

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

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



					def    A__		(			self     ):
							import bitsandbytes as bnb

							from transformers import TaForConditionalGeneration

							# test with `t5-small`
							UpperCAmelCase_						=				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_						=				self.tokenizer(self.input_text    ,	return_tensors="pt"     ).to(0     )
							UpperCAmelCase_						=				model.generate(**lowerCAmelCase     )

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





class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					def    A__		(			self     ):
							super().setUp()
							# model_name
							UpperCAmelCase_						=				"bigscience/bloom-560m"
							UpperCAmelCase_						=				"t5-small"

							# Different types of model

							UpperCAmelCase_						=				AutoModel.from_pretrained(self.model_name    ,	load_in_abit=lowerCAmelCase    ,	device_map="auto"     )
							# Sequence classification model
							UpperCAmelCase_						=				AutoModelForSequenceClassification.from_pretrained(
							    self.model_name    ,	load_in_abit=lowerCAmelCase    ,	device_map="auto"     )
							# CausalLM model
							UpperCAmelCase_						=				AutoModelForCausalLM.from_pretrained(self.model_name    ,	load_in_abit=lowerCAmelCase    ,	device_map="auto"     )
							# Seq2seq model
							UpperCAmelCase_						=				AutoModelForSeqaSeqLM.from_pretrained(
							    self.seq_to_seq_name    ,	load_in_abit=lowerCAmelCase    ,	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 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







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



					def    A__		(			self     ):
							del self.pipe

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



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





@require_torch_multi_gpu
class 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







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



					def    A__		(			self     ):
							UpperCAmelCase_						=				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_						=				self.tokenizer(self.input_text    ,	return_tensors="pt"     )

							# Second real batch
							UpperCAmelCase_						=				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 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					def    A__		(			self     ):
							UpperCAmelCase_						=				"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
							UpperCAmelCase_						=				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_						=				False  # freeze the model - train adapters later
									if param.ndim == 1:
											# cast the small parameters (e.g. layernorm) to fp32 for stability
											UpperCAmelCase_						=				param.data.to(torch.floataa     )

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

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

							# Step 4: Check if the gradient is not None
							with torch.cuda.amp.autocast():
									UpperCAmelCase_						=				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 	lowerCamelCase      (						lowercase__  ):







					'''simple docstring'''







					lowerCAmelCase_				:    Dict									=  'gpt2-xl'
					lowerCAmelCase_				:    int									=  3.3_1_9_1_8_5_4_8_5_4_1_5_2_1_8_7






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE     				=						{"configuration_xlnet": ["XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLNetConfig"]}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						["XLNetTokenizer"]

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						["XLNetTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						[
			    "XLNET_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "XLNetForMultipleChoice",
			    "XLNetForQuestionAnswering",
			    "XLNetForQuestionAnsweringSimple",
			    "XLNetForSequenceClassification",
			    "XLNetForTokenClassification",
			    "XLNetLMHeadModel",
			    "XLNetModel",
			    "XLNetPreTrainedModel",
			    "load_tf_weights_in_xlnet",
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			SCREAMING_SNAKE_CASE     				=						[
			    "TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "TFXLNetForMultipleChoice",
			    "TFXLNetForQuestionAnsweringSimple",
			    "TFXLNetForSequenceClassification",
			    "TFXLNetForTokenClassification",
			    "TFXLNetLMHeadModel",
			    "TFXLNetMainLayer",
			    "TFXLNetModel",
			    "TFXLNetPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_xlnet import XLNetTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_xlnet_fast import XLNetTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_xlnet import (
						    XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
						    XLNetForMultipleChoice,
						    XLNetForQuestionAnswering,
						    XLNetForQuestionAnsweringSimple,
						    XLNetForSequenceClassification,
						    XLNetForTokenClassification,
						    XLNetLMHeadModel,
						    XLNetModel,
						    XLNetPreTrainedModel,
						    load_tf_weights_in_xlnet,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_xlnet import (
						    TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFXLNetForMultipleChoice,
						    TFXLNetForQuestionAnsweringSimple,
						    TFXLNetForSequenceClassification,
						    TFXLNetForTokenClassification,
						    TFXLNetLMHeadModel,
						    TFXLNetMainLayer,
						    TFXLNetModel,
						    TFXLNetPreTrainedModel,
						)

else:
			import sys

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

