


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




										"""simple docstring"""


										def __init__(							self :						Dict					,			lowerCamelCase :						nn.Module					,			lowerCamelCase :						int       )			->					List[str]:
																	super().__init__()
																	__snake_case	:    Union[str, Any]											=     module
																	__snake_case	:    Tuple											=     nn.Sequential(
																	    nn.Linear(module.in_features					,			lowerCamelCase					,			bias=lowerCamelCase       )					,			nn.Linear(lowerCamelCase					,			module.out_features					,			bias=lowerCamelCase       )					,			)
																	__snake_case	:    Union[str, Any]											=     (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 					__snake_case      (							self :						str					,			lowerCamelCase :						Tuple					,			*lowerCamelCase :						Optional[int]					,			**lowerCamelCase :						Dict       )			->					Dict:
																	return self.module(lowerCamelCase					,			*lowerCamelCase					,			**lowerCamelCase       ) + self.adapter(lowerCamelCase       )

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




							"""simple docstring"""


							__UpperCAmelCase					:							List[str]  						=						"bigscience/bloom-1b7"

							# Constant values
							__UpperCAmelCase					:							str  						=						2.109659552692574

							__UpperCAmelCase					:							Optional[Any]  						=						"Hello my name is"
							__UpperCAmelCase					:							List[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"   )
							__UpperCAmelCase					:							Optional[Any]  						=						10

							def 					__snake_case      (							self :						Tuple       )			->					Optional[Any]:
														# Models and tokenizer
														__snake_case	:    Any											=     AutoTokenizer.from_pretrained(self.model_name       )

class        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						List[str]       )			->					Optional[Any]:
														super().setUp()

														# Models and tokenizer
														__snake_case	:    str											=     AutoModelForCausalLM.from_pretrained(
														    self.model_name					,			torch_dtype=torch.floataa					,			device_map="auto"       )
														__snake_case	:    Optional[int]											=     AutoModelForCausalLM.from_pretrained(self.model_name					,			load_in_abit=lowerCamelCase					,			device_map="auto"       )

							def 					__snake_case      (							self :						Optional[Any]       )			->					Dict:
														del self.model_fpaa
														del self.model_abit

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

							def 					__snake_case      (							self :						Optional[Any]       )			->					Any:
														__snake_case	:    List[Any]											=     self.model_abit.config

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

														__snake_case	:    List[str]											=     config.to_dict()
														__snake_case	:    int											=     config.to_diff_dict()

														__snake_case	:    Any											=     config.to_json_string()

							def 					__snake_case      (							self :						int       )			->					Tuple:
														from bitsandbytes.nn import Paramsabit

														__snake_case	:    Tuple											=     self.model_fpaa.get_memory_footprint()
														__snake_case	:    Optional[int]											=     self.model_abit.get_memory_footprint()

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

							def 					__snake_case      (							self :						str       )			->					Optional[Any]:
														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 					__snake_case      (							self :						int       )			->					int:
														__snake_case	:    str											=     self.tokenizer(self.input_text					,			return_tensors="pt"       )
														__snake_case	:    str											=     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 					__snake_case      (							self :						List[str]       )			->					Tuple:
														__snake_case	:    str											=     BitsAndBytesConfig()
														__snake_case	:    Dict											=     True

														__snake_case	:    int											=     AutoModelForCausalLM.from_pretrained(
														    self.model_name					,			quantization_config=lowerCamelCase					,			device_map="auto"       )

														__snake_case	:    Tuple											=     self.tokenizer(self.input_text					,			return_tensors="pt"       )
														__snake_case	:    List[str]											=     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 					__snake_case      (							self :						Tuple       )			->					Tuple:
														with self.assertRaises(lowerCamelCase       ), tempfile.TemporaryDirectory() as tmpdirname:
																					self.model_abit.save_pretrained(lowerCamelCase       )

							def 					__snake_case      (							self :						str       )			->					Union[str, Any]:
														__snake_case	:    Tuple											=     BitsAndBytesConfig()

														with self.assertRaises(lowerCamelCase       ):
																					__snake_case	:    Tuple											=     AutoModelForCausalLM.from_pretrained(
																					    self.model_name					,			quantization_config=lowerCamelCase					,			load_in_abit=lowerCamelCase					,			device_map="auto"					,			bnb_abit_quant_type="nf4"					,			)

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

														__snake_case	:    str											=     self.model_fpaa.to(torch.floataa       )
														__snake_case	:    Dict											=     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	:    Optional[int]											=     self.model_fpaa.to("cpu"       )

														# Check this does not throw an error
														__snake_case	:    str											=     self.model_fpaa.half()

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






							def 					__snake_case      (							self :						Dict       )			->					Optional[int]:
														__snake_case	:    List[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        a	(unittest.TestCase   ):




							"""simple docstring"""


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

							def 					__snake_case      (							self :						List[Any]       )			->					Dict:
														gc.collect()
														torch.cuda.empty_cache()

							def 					__snake_case      (							self :						Any       )			->					List[str]:
														from transformers import TaForConditionalGeneration

														__snake_case	:    int											=     TaForConditionalGeneration._keep_in_fpaa_modules
														__snake_case	:    Any											=     None

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

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






							def 					__snake_case      (							self :						Optional[Any]       )			->					str:
														import bitsandbytes as bnb

														from transformers import TaForConditionalGeneration

														# test with `t5-small`
														__snake_case	:    List[str]											=     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       )       )

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

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

class        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						Union[str, Any]       )			->					Dict:
														super().setUp()
														# model_name
														__snake_case	:    Any											=     "bigscience/bloom-560m"
														__snake_case	:    int											=     "t5-small"

														# Different types of model

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

							def 					__snake_case      (							self :						Union[str, Any]       )			->					Optional[int]:
														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 					__snake_case      (							self :						int       )			->					int:
														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        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						int       )			->					str:
														super().setUp()

							def 					__snake_case      (							self :						int       )			->					Optional[int]:
														del self.pipe

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






							def 					__snake_case      (							self :						Tuple       )			->					List[str]:
														__snake_case	:    Optional[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
														__snake_case	:    Optional[Any]											=     self.pipe(self.input_text       )
														self.assertIn(pipeline_output[0]["generated_text"]					,			self.EXPECTED_OUTPUTS       )

@require_torch_multi_gpu
class        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						List[str]       )			->					str:
														super().setUp()






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

														# Second real batch
														__snake_case	:    Any											=     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        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							def 					__snake_case      (							self :						Dict       )			->					Any:
														__snake_case	:    Dict											=     "facebook/opt-350m"
														super().setUp()






							def 					__snake_case      (							self :						List[str]       )			->					Tuple:
														if version.parse(importlib.metadata.version("bitsandbytes"       )       ) < version.parse("0.37.0"       ):
																					return

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

        # Step 2: add adapters
														for _, module in model.named_modules():
																					if "OPTAttention" in repr(type(lowerCamelCase       )       ):
																												__snake_case	:    Optional[int]											=     LoRALayer(module.q_proj					,			rank=16       )
																												__snake_case	:    Any											=     LoRALayer(module.k_proj					,			rank=16       )
																												__snake_case	:    Optional[int]											=     LoRALayer(module.v_proj					,			rank=16       )

        # Step 3: dummy batch
														__snake_case	:    Dict											=     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	:    Optional[Any]											=     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        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							__UpperCAmelCase					:							List[str]  						=						"gpt2-xl"
							__UpperCAmelCase					:							Any  						=						3.3191854854152187





import warnings

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


_snake_case  :	Optional[Any]			      =       logging.get_logger(__name__)

_snake_case  :	str			      =       {
    "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/config.json",
    "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/config.json",
}

class        a	(_lowerCAmelCase   ):




							"""simple docstring"""


							__UpperCAmelCase					:							List[Any]  						=						"xlnet"
							__UpperCAmelCase					:							Dict  						=						["mems"]
							__UpperCAmelCase					:							Tuple  						=						{
							    "n_token": "vocab_size",  # Backward compatibility
							    "hidden_size": "d_model",
							    "num_attention_heads": "n_head",
							    "num_hidden_layers": "n_layer",
							}

							def __init__(							self :						List[Any]					,			lowerCamelCase :						Optional[int]=32000					,			lowerCamelCase :						str=1024					,			lowerCamelCase :						int=24					,			lowerCamelCase :						List[Any]=16					,			lowerCamelCase :						Any=4096					,			lowerCamelCase :						Optional[Any]="gelu"					,			lowerCamelCase :						Optional[Any]=True					,			lowerCamelCase :						Optional[int]="bi"					,			lowerCamelCase :						Union[str, Any]=0.02					,			lowerCamelCase :						Union[str, Any]=1E-12					,			lowerCamelCase :						Dict=0.1					,			lowerCamelCase :						int=512					,			lowerCamelCase :						str=None					,			lowerCamelCase :						Dict=True					,			lowerCamelCase :						str=False					,			lowerCamelCase :						Any=False					,			lowerCamelCase :						Optional[Any]=-1					,			lowerCamelCase :						List[Any]=False					,			lowerCamelCase :						Optional[int]="last"					,			lowerCamelCase :						Dict=True					,			lowerCamelCase :						Union[str, Any]="tanh"					,			lowerCamelCase :						int=0.1					,			lowerCamelCase :						Any=5					,			lowerCamelCase :						List[str]=5					,			lowerCamelCase :						List[Any]=5					,			lowerCamelCase :						Tuple=1					,			lowerCamelCase :						Dict=2					,			**lowerCamelCase :						List[str]					,			)			->					List[Any]:
														__snake_case	:    Optional[Any]											=     vocab_size
														__snake_case	:    Any											=     d_model
														__snake_case	:    int											=     n_layer
														__snake_case	:    Optional[int]											=     n_head
														if d_model % n_head != 0:
																					raise ValueError(F'\'d_model % n_head\' ({d_model % n_head}) should be equal to 0'       )
														if "d_head" in kwargs:
																					if kwargs["d_head"] != d_model // n_head:
																												raise ValueError(
																												    F'`d_head` ({kwargs["d_head"]}) should be equal to `d_model // n_head` ({d_model // n_head})'       )
														__snake_case	:    Any											=     d_model // n_head
														__snake_case	:    Tuple											=     ff_activation
														__snake_case	:    int											=     d_inner
														__snake_case	:    Optional[Any]											=     untie_r
														__snake_case	:    Optional[Any]											=     attn_type

														__snake_case	:    List[str]											=     initializer_range
														__snake_case	:    Tuple											=     layer_norm_eps

														__snake_case	:    str											=     dropout
														__snake_case	:    Any											=     mem_len
														__snake_case	:    Tuple											=     reuse_len
														__snake_case	:    Union[str, Any]											=     bi_data
														__snake_case	:    List[Any]											=     clamp_len
														__snake_case	:    Optional[int]											=     same_length

														__snake_case	:    List[str]											=     summary_type
														__snake_case	:    Any											=     summary_use_proj
														__snake_case	:    Any											=     summary_activation
														__snake_case	:    Optional[int]											=     summary_last_dropout
														__snake_case	:    Dict											=     start_n_top
														__snake_case	:    Optional[Any]											=     end_n_top

														__snake_case	:    List[str]											=     bos_token_id
														__snake_case	:    Any											=     pad_token_id
														__snake_case	:    Optional[Any]											=     eos_token_id

														if "use_cache" in kwargs:
																					warnings.warn(
																					    "The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`"
																					    " instead."					,			lowerCamelCase					,			)
																					__snake_case	:    Any											=     kwargs["use_cache"]

														__snake_case	:    int											=     use_mems_eval
														__snake_case	:    Optional[int]											=     use_mems_train
														super().__init__(pad_token_id=lowerCamelCase					,			bos_token_id=lowerCamelCase					,			eos_token_id=lowerCamelCase					,			**lowerCamelCase       )

							@property
							def 					__snake_case      (							self :						Any       )			->					Optional[int]:
														logger.info(F'The model {self.model_type} is one of the few models that has no sequence length limit.'       )
														return -1






							@max_position_embeddings.setter
							def 					__snake_case      (							self :						Union[str, Any]					,			lowerCamelCase :						Any       )			->					Optional[Any]:
														# Message copied from Transformer-XL documentation
														raise NotImplementedError(
														    F'The model {self.model_type} is one of the few models that has no sequence length limit.'       )

