



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  UpperCAmelCase__    (    lowercase__   ) ->      Optional[int]:
						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	:       List[Any]			,     lowercase	:       nn.Module			,     lowercase	:       int   )						->					Union[str, Any]:
												"""simple docstring"""







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







						def 		snake_case__			(				self	:       Any			,     lowercase	:       List[str]			,     *lowercase	:       List[Any]			,     **lowercase	:       Union[str, Any]   )						->					Optional[int]:
												"""simple docstring"""







												return self.module(lowercase			,     *lowercase			,     **lowercase   ) + self.adapter(lowercase   )
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  Tuple      =			"""bigscience/bloom-1b7"""

			# Constant values
			lowercase__      :  int      =			2.1_09_65_95_52_69_25_74

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





			def 		snake_case__			(				self	:       Optional[int]   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				AutoTokenizer.from_pretrained(self.model_name   )





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Union[str, Any]   )						->					List[str]:
									"""simple docstring"""







									super().setUp()

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





			def 		snake_case__			(				self	:       List[str]   )						->					Union[str, Any]:
									"""simple docstring"""







									del self.model_fpaa
									del self.model_abit

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





			def 		snake_case__			(				self	:       Any   )						->					Optional[Any]:
									"""simple docstring"""







									__lowercase					=				self.model_abit.config

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

									__lowercase					=				config.to_dict()
									__lowercase					=				config.to_diff_dict()

									__lowercase					=				config.to_json_string()





			def 		snake_case__			(				self	:       int   )						->					List[str]:
									"""simple docstring"""







									from bitsandbytes.nn import Paramsabit

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

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





			def 		snake_case__			(				self	:       str   )						->					str:
									"""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(lowercase			,     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	:       Optional[int]   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				self.tokenizer(self.input_text			,     return_tensors="""pt"""   )
									__lowercase					=				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=lowercase   )			,     self.EXPECTED_OUTPUTS   )





			def 		snake_case__			(				self	:       List[Any]   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				BitsAndBytesConfig()
									__lowercase					=				True

									__lowercase					=				AutoModelForCausalLM.from_pretrained(
									    self.model_name			,     quantization_config=lowercase			,     device_map="""auto"""   )

									__lowercase					=				self.tokenizer(self.input_text			,     return_tensors="""pt"""   )
									__lowercase					=				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=lowercase   )			,     self.EXPECTED_OUTPUTS   )





			def 		snake_case__			(				self	:       Tuple   )						->					int:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       str   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				BitsAndBytesConfig()

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





			def 		snake_case__			(				self	:       Any   )						->					int:
									"""simple docstring"""







									with self.assertRaises(lowercase   ):
															# Tries with `str`
															self.model_abit.to("""cpu"""   )

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

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

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

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

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

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

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

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

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







			def 		snake_case__			(				self	:       Optional[Any]   )						->					List[Any]:
									"""simple docstring"""







									__lowercase					=				AutoModelForSeqaSeqLM.from_pretrained("""t5-small"""			,     load_in_abit=lowercase			,     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 		snake_case__			(				cls	:       Optional[Any]   )						->					List[Any]:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       Any   )						->					Optional[Any]:
									"""simple docstring"""







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





			def 		snake_case__			(				self	:       List[str]   )						->					List[Any]:
									"""simple docstring"""







									from transformers import TaForConditionalGeneration

									__lowercase					=				TaForConditionalGeneration._keep_in_fpaa_modules
									__lowercase					=				None

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

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







			def 		snake_case__			(				self	:       Dict   )						->					List[str]:
									"""simple docstring"""







									import bitsandbytes as bnb

									from transformers import TaForConditionalGeneration

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

									__lowercase					=				self.tokenizer(self.input_text			,     return_tensors="""pt"""   ).to(0   )
									__lowercase					=				model.generate(**lowercase   )

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





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Optional[Any]   )						->					List[str]:
									"""simple docstring"""







									super().setUp()
									# model_name
									__lowercase					=				"""bigscience/bloom-560m"""
									__lowercase					=				"""t5-small"""

									# Different types of model

									__lowercase					=				AutoModel.from_pretrained(self.model_name			,     load_in_abit=lowercase			,     device_map="""auto"""   )
									# Sequence classification model
									__lowercase					=				AutoModelForSequenceClassification.from_pretrained(
									    self.model_name			,     load_in_abit=lowercase			,     device_map="""auto"""   )
									# CausalLM model
									__lowercase					=				AutoModelForCausalLM.from_pretrained(self.model_name			,     load_in_abit=lowercase			,     device_map="""auto"""   )
									# Seq2seq model
									__lowercase					=				AutoModelForSeqaSeqLM.from_pretrained(
									    self.seq_to_seq_name			,     load_in_abit=lowercase			,     device_map="""auto"""   )





			def 		snake_case__			(				self	:       Optional[int]   )						->					Tuple:
									"""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 		snake_case__			(				self	:       Optional[Any]   )						->					Any:
									"""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 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       int   )						->					str:
									"""simple docstring"""







									super().setUp()





			def 		snake_case__			(				self	:       Optional[Any]   )						->					Tuple:
									"""simple docstring"""







									del self.pipe

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







			def 		snake_case__			(				self	:       int   )						->					Optional[int]:
									"""simple docstring"""







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





@require_torch_multi_gpu
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Tuple   )						->					Tuple:
									"""simple docstring"""







									super().setUp()







			def 		snake_case__			(				self	:       int   )						->					Dict:
									"""simple docstring"""







									__lowercase					=				AutoModelForCausalLM.from_pretrained(
									    self.model_name			,     load_in_abit=lowercase			,     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
									__lowercase					=				self.tokenizer(self.input_text			,     return_tensors="""pt"""   )

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





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Optional[int]   )						->					Optional[Any]:
									"""simple docstring"""







									__lowercase					=				"""facebook/opt-350m"""
									super().setUp()







			def 		snake_case__			(				self	:       Dict   )						->					str:
									"""simple docstring"""







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

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

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

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

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

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

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

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





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  List[Any]      =			"""gpt2-xl"""
			lowercase__      :  Union[str, Any]      =			3.31_91_85_48_54_15_21_87










import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


UpperCamelCase__ =						"scheduler_config.json"
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  List[Any]      =			1
			lowercase__      :  Tuple      =			2
			lowercase__      :  Union[str, Any]      =			3
			lowercase__      :  Union[str, Any]      =			4
			lowercase__      :  str      =			5
			lowercase__      :  Any      =			6
			lowercase__      :  Any      =			7
			lowercase__      :  List[str]      =			8
			lowercase__      :  Union[str, Any]      =			9
			lowercase__      :  int      =			10
			lowercase__      :  List[str]      =			11
			lowercase__      :  List[Any]      =			12
			lowercase__      :  str      =			13
			lowercase__      :  Optional[int]      =			14





@dataclass
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  torch.FloatTensor





class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  Optional[int]      =			SCHEDULER_CONFIG_NAME
			lowercase__      :  int      =			[]
			lowercase__      :  Dict      =			True





			@classmethod
			def 		snake_case__			(				cls	:       str			,     lowercase	:       Dict[str, Any] = None			,     lowercase	:       Optional[str] = None			,     lowercase	:       Any=False			,     **lowercase	:       List[str]			,     )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase	,			__lowercase	,			__lowercase					=				cls.load_config(
									    pretrained_model_name_or_path=lowercase			,     subfolder=lowercase			,     return_unused_kwargs=lowercase			,     return_commit_hash=lowercase			,     **lowercase			,     )
									return cls.from_config(lowercase			,     return_unused_kwargs=lowercase			,     **lowercase   )





			def 		snake_case__			(				self	:       Dict			,     lowercase	:       Union[str, os.PathLike]			,     lowercase	:       bool = False			,     **lowercase	:       List[str]   )						->					Optional[Any]:
									"""simple docstring"""







									self.save_config(save_directory=lowercase			,     push_to_hub=lowercase			,     **lowercase   )





			@property
			def 		snake_case__			(				self	:       Optional[Any]   )						->					Dict:
									"""simple docstring"""







									return self._get_compatibles()







			@classmethod
			def 		snake_case__			(				cls	:       List[str]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				list(set([cls.__name__] + cls._compatibles   )   )
									__lowercase					=				importlib.import_module(__name__.split("""."""   )[0]   )
									__lowercase					=				[
									    getattr(lowercase			,     lowercase   ) for c in compatible_classes_str if hasattr(lowercase			,     lowercase   )
									]
									return compatible_classes





