
"""simple docstring"""


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__      (				UpperCAmelCase_					)->				Optional[int]:







       """simple docstring"""
       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						):






             def __init__(	self			:				List[Any]							,   UpperCAmelCase_			:				nn.Module							,   UpperCAmelCase_			:				int      )->	str:







                    """simple docstring"""
                    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 						_SCREAMING_SNAKE_CASE							(	self			:				Any							,   UpperCAmelCase_			:				Union[str, Any]							,   *UpperCAmelCase_			:				Any							,   **UpperCAmelCase_			:				Union[str, Any]      )->	List[Any]:







                    """simple docstring"""
                    return self.module(UpperCAmelCase_							,   *UpperCAmelCase_							,   **UpperCAmelCase_      ) + self.adapter(UpperCAmelCase_      )



@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 	__a	(   unittest.TestCase						):
       # We keep the constants inside the init function and model loading inside setUp function

       # We need to test on relatively large models (aka >1b parameters otherwise the quantiztion may not work as expected)
       # Therefore here we use only bloom-1b3 to test our module
       UpperCamelCase_		:       List[str]				=      '''bigscience/bloom-1b7'''

       # Constant values
       UpperCamelCase_		:       int				=      2.1_0_9_6_5_9_5_5_2_6_9_2_5_7_4

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






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	Dict:







              """simple docstring"""
              # Models and tokenizer
              UpperCamelCase       			=		AutoTokenizer.from_pretrained(self.model_name      )





class 	__a	(   _lowerCAmelCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]      )->	List[str]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Dict:







              """simple docstring"""
              del self.model_fpaa
              del self.model_abit

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






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	Optional[int]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	Union[str, Any]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]      )->	Dict:







              """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(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 						_SCREAMING_SNAKE_CASE							(	self			:				str      )->	str:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	Union[str, Any]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	int:







              """simple docstring"""
              with self.assertRaises(UpperCAmelCase_      ), tempfile.TemporaryDirectory() as tmpdirname:
                     self.model_abit.save_pretrained(UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	Optional[Any]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple      )->	Tuple:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	List[Any]:







              """simple docstring"""
              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						):






       @classmethod
       def 						_SCREAMING_SNAKE_CASE							(	cls			:				Dict      )->	str:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	Optional[int]:







              """simple docstring"""
              gc.collect()
              torch.cuda.empty_cache()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]      )->	Optional[Any]:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	List[Any]:







              """simple docstring"""
              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	(   _lowerCAmelCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	str:







              """simple docstring"""
              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 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	List[str]:







              """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 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple      )->	int:







              """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 	__a	(   _lowerCAmelCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]      )->	List[Any]:







              """simple docstring"""
              super().setUp()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]      )->	Dict:







              """simple docstring"""
              del self.pipe

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






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]      )->	Tuple:







              """simple docstring"""
              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	(   _lowerCAmelCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	Any:







              """simple docstring"""
              super().setUp()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	List[str]:







              """simple docstring"""
              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	(   _lowerCAmelCase						):






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		"facebook/opt-350m"
              super().setUp()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	List[str]:







              """simple docstring"""
              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	(   _lowerCAmelCase						):
       UpperCamelCase_		:       List[Any]				=      '''gpt2-xl'''
       UpperCamelCase_		:       Any				=      3.3_1_9_1_8_5_4_8_5_4_1_5_2_1_8_7







"""simple docstring"""


import warnings

from ...utils import logging
from .image_processing_yolos import YolosImageProcessor


SCREAMING_SNAKE_CASE			  =  logging.get_logger(__name__)



class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				str							,   *UpperCAmelCase_			:				str							,   **UpperCAmelCase_			:				Any      )->	None:







              """simple docstring"""
              warnings.warn(
                  "The class YolosFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
                  " use YolosImageProcessor instead."							,   UpperCAmelCase_							,   )
              super().__init__(*UpperCAmelCase_							,   **UpperCAmelCase_      )





