






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     (__lowercase):
   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		,		_UpperCAmelCase		,		_UpperCAmelCase			)      ->			List[Any]:
      super().__init__()
      UpperCamelCase_						=  module
      UpperCamelCase_						=  nn.Sequential(
          nn.Linear(module.in_features		,		_UpperCAmelCase		,		bias=_UpperCAmelCase			)		,		nn.Linear(_UpperCAmelCase		,		module.out_features		,		bias=_UpperCAmelCase			)		,		)
      UpperCamelCase_						=  (2.0 / (5 * min(module.in_features		,		module.out_features			))) ** 0.5
      nn.init.normal_(self.adapter[0].weight		,		std=_UpperCAmelCase			)
      nn.init.zeros_(self.adapter[1].weight			)
      self.adapter.to(module.weight.device			)




   def  _UpperCAmelCase  (  self		,		_UpperCAmelCase		,		*_UpperCAmelCase		,		**_UpperCAmelCase			)      ->			Optional[int]:
      return self.module(_UpperCAmelCase		,		*_UpperCAmelCase		,		**_UpperCAmelCase			) + self.adapter(_UpperCAmelCase			)


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






  """simple docstring"""





  A_						       = """bigscience/bloom-1b7"""

  # Constant values
  A_						       = 2.109_659_552_692_574

  A_						       = """Hello my name is"""
  A_						       = 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"""     )
  A_						       = 10




  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     # Models and tokenizer
     UpperCamelCase_						=  AutoTokenizer.from_pretrained(self.model_name			)




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     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=_UpperCAmelCase		,		device_map='auto'			)




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     del self.model_fpaa
     del self.model_abit

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




  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     UpperCamelCase_						=  self.model_abit.config

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

     UpperCamelCase_						=  config.to_dict()
     UpperCamelCase_						=  config.to_diff_dict()

     UpperCamelCase_						=  config.to_json_string()




  def  _UpperCAmelCase  (  self			)      ->			int:
     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  _UpperCAmelCase  (  self			)      ->			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(_UpperCAmelCase		,		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			)      ->			Dict:
     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=_UpperCAmelCase			)		,		self.EXPECTED_OUTPUTS			)




  def  _UpperCAmelCase  (  self			)      ->			Tuple:
     UpperCamelCase_						=  BitsAndBytesConfig()
     UpperCamelCase_						=  True

     UpperCamelCase_						=  AutoModelForCausalLM.from_pretrained(
         self.model_name		,		quantization_config=_UpperCAmelCase		,		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=_UpperCAmelCase			)		,		self.EXPECTED_OUTPUTS			)




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




  def  _UpperCAmelCase  (  self			)      ->			List[Any]:
     UpperCamelCase_						=  BitsAndBytesConfig()

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




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     with self.assertRaises(_UpperCAmelCase			):
        # Tries with `str`
        self.model_abit.to('cpu'			)

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

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

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

     with self.assertRaises(_UpperCAmelCase			):
        # 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  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  AutoModelForSeqaSeqLM.from_pretrained('t5-small'		,		load_in_abit=_UpperCAmelCase		,		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  _UpperCAmelCase  (  cls			)      ->			Tuple:
     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  _UpperCAmelCase  (  self			)      ->			List[Any]:
     gc.collect()
     torch.cuda.empty_cache()




  def  _UpperCAmelCase  (  self			)      ->			Any:
     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=_UpperCAmelCase		,		device_map='auto'			)
     UpperCamelCase_						=  self.tokenizer(self.input_text		,		return_tensors='pt'			).to(0			)
     UpperCamelCase_						=  model.generate(**_UpperCAmelCase			)

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




  def  _UpperCAmelCase  (  self			)      ->			str:
     import bitsandbytes as bnb

     from transformers import TaForConditionalGeneration

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

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




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     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=_UpperCAmelCase		,		device_map='auto'			)
     # Sequence classification model
     UpperCamelCase_						=  AutoModelForSequenceClassification.from_pretrained(
         self.model_name		,		load_in_abit=_UpperCAmelCase		,		device_map='auto'			)
     # CausalLM model
     UpperCamelCase_						=  AutoModelForCausalLM.from_pretrained(self.model_name		,		load_in_abit=_UpperCAmelCase		,		device_map='auto'			)
     # Seq2seq model
     UpperCamelCase_						=  AutoModelForSeqaSeqLM.from_pretrained(
         self.seq_to_seq_name		,		load_in_abit=_UpperCAmelCase		,		device_map='auto'			)




  def  _UpperCAmelCase  (  self			)      ->			Union[str, Any]:
     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			)      ->			Union[str, Any]:
     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							(	UpperCAmelCase__     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Any:
     super().setUp()




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     del self.pipe

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




  def  _UpperCAmelCase  (  self			)      ->			Any:
     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     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			Dict:
     super().setUp()




  def  _UpperCAmelCase  (  self			)      ->			int:
     UpperCamelCase_						=  AutoModelForCausalLM.from_pretrained(
         self.model_name		,		load_in_abit=_UpperCAmelCase		,		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=_UpperCAmelCase			)		,		self.EXPECTED_OUTPUTS			)




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  def  _UpperCAmelCase  (  self			)      ->			str:
     UpperCamelCase_						=  'facebook/opt-350m'
     super().setUp()




  def  _UpperCAmelCase  (  self			)      ->			Optional[Any]:
     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=_UpperCAmelCase			)

     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(_UpperCAmelCase			)			):
           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(**_UpperCAmelCase			)
        out.logits.norm().backward()

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




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  A_						       = """gpt2-xl"""
  A_						       = 3.3_191_854_854_152_187













import copy
from typing import TYPE_CHECKING, Any, Mapping, Optional, OrderedDict

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


if TYPE_CHECKING:
 from ... import PreTrainedTokenizerBase, TensorType

snake_case__   :						List[str]													=   logging.get_logger(__name__)


class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  A_						       = """vision-encoder-decoder"""
  A_						       = True




  def __init__(  self		,		**_UpperCAmelCase			)      ->			Dict:
     super().__init__(**_UpperCAmelCase			)
     if "encoder" not in kwargs or "decoder" not in kwargs:
        raise ValueError(
            f"""A configuraton of type {self.model_type} cannot be instantiated because """
            f"""not both `encoder` and `decoder` sub-configurations are passed, but only {kwargs}"""			)

     UpperCamelCase_						=  kwargs.pop('encoder'			)
     UpperCamelCase_						=  encoder_config.pop('model_type'			)
     UpperCamelCase_						=  kwargs.pop('decoder'			)
     UpperCamelCase_						=  decoder_config.pop('model_type'			)

     UpperCamelCase_						=  AutoConfig.for_model(_UpperCAmelCase		,		**_UpperCAmelCase			)
     UpperCamelCase_						=  AutoConfig.for_model(_UpperCAmelCase		,		**_UpperCAmelCase			)
     UpperCamelCase_						=  True




  @classmethod
  def  _UpperCAmelCase  (  cls		,		_UpperCAmelCase		,		_UpperCAmelCase		,		**_UpperCAmelCase			)      ->			PretrainedConfig:
     logger.info('Setting `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'			)
     UpperCamelCase_						=  True
     UpperCamelCase_						=  True

     return cls(encoder=encoder_config.to_dict()		,		decoder=decoder_config.to_dict()		,		**_UpperCAmelCase			)




  def  _UpperCAmelCase  (  self			)      ->			int:
     UpperCamelCase_						=  copy.deepcopy(self.__dict__			)
     UpperCamelCase_						=  self.encoder.to_dict()
     UpperCamelCase_						=  self.decoder.to_dict()
     UpperCamelCase_						=  self.__class__.model_type
     return output




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  A_						       = version.parse("""1.11"""     )




  @property
  def  _UpperCAmelCase  (  self			)      ->			Mapping[str, Mapping[int, str]]:
     return OrderedDict(
         [
             ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
         ]			)




  @property
  def  _UpperCAmelCase  (  self			)      ->			float:
     return 1e-4




  @property
  def  _UpperCAmelCase  (  self			)      ->			Mapping[str, Mapping[int, str]]:
     return OrderedDict({'last_hidden_state': {0: 'batch', 1: 'encoder_sequence'}}			)




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  @property
  def  _UpperCAmelCase  (  self			)      ->			Mapping[str, Mapping[int, str]]:
     UpperCamelCase_						=  OrderedDict()
     UpperCamelCase_						=  {0: 'batch', 1: 'past_decoder_sequence + sequence'}
     UpperCamelCase_						=  {0: 'batch', 1: 'past_decoder_sequence + sequence'}
     UpperCamelCase_						=  {0: 'batch', 1: 'encoder_sequence'}

     return common_inputs




  def  _UpperCAmelCase  (  self		,		_UpperCAmelCase		,		_UpperCAmelCase = -1		,		_UpperCAmelCase = -1		,		_UpperCAmelCase = False		,		_UpperCAmelCase = None		,		)      ->			Mapping[str, Any]:
     import torch

     UpperCamelCase_						=  OrderedDict()

     UpperCamelCase_						=  super().generate_dummy_inputs(
         _UpperCAmelCase		,		batch_size=_UpperCAmelCase		,		seq_length=_UpperCAmelCase		,		is_pair=_UpperCAmelCase		,		framework=_UpperCAmelCase			)

     UpperCamelCase_		,   UpperCamelCase_						=  dummy_input['input_ids'].shape
     UpperCamelCase_						=  (batch, encoder_sequence, self._config.encoder_hidden_size)
     UpperCamelCase_						=  dummy_input.pop('input_ids'			)
     UpperCamelCase_						=  dummy_input.pop('attention_mask'			)
     UpperCamelCase_						=  torch.zeros(_UpperCAmelCase			)

     return common_inputs




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  @property
  def  _UpperCAmelCase  (  self			)      ->			None:
     pass




  def  _UpperCAmelCase  (  self		,		_UpperCAmelCase			)      ->			OnnxConfig:
     return VisionEncoderDecoderEncoderOnnxConfig(_UpperCAmelCase			)




  def  _UpperCAmelCase  (  self		,		_UpperCAmelCase		,		_UpperCAmelCase		,		_UpperCAmelCase = "default"			)      ->			OnnxConfig:
     UpperCamelCase_						=  encoder_config.hidden_size
     return VisionEncoderDecoderDecoderOnnxConfig(_UpperCAmelCase		,		_UpperCAmelCase			)





