"""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     UpperCAmelCase__ (       lowerCAmelCase__  :Optional[Any]	)				->					Optional[Any]:
      '''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	, __lowerCAmelCase	, __lowerCAmelCase		):







                   """simple docstring"""



                   super().__init__()
                   lowercase														=			module
                   lowercase														=			nn.Sequential(
                       nn.Linear(module.in_features	, __lowerCAmelCase	, bias=__lowerCAmelCase		)	, nn.Linear(__lowerCAmelCase	, module.out_features	, bias=__lowerCAmelCase		)	, )
                   lowercase														=			(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		):







                   """simple docstring"""



                   return self.module(__lowerCAmelCase	, *__lowerCAmelCase	, **__lowerCAmelCase		) + self.adapter(__lowerCAmelCase		)





@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
      snake_case__    :     Optional[Any]													=			'bigscience/bloom-1b7'

      # Constant values
      snake_case__    :     Optional[Any]													=			2.1_09_65_95_52_69_25_74

      snake_case__    :     Optional[int]													=			'Hello my name is'
      snake_case__    :     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'			)
      snake_case__    :     Dict													=			10

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			AutoTokenizer.from_pretrained(self.model_name		)






class 				_A		( lowerCAmelCase			):

      def        A__ (		self		):







            """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=__lowerCAmelCase	, device_map="""auto"""		)

      def        A__ (		self		):







            """simple docstring"""



            del self.model_fpaa
            del self.model_abit

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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_abit.config

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

            lowercase														=			config.to_dict()
            lowercase														=			config.to_diff_dict()

            lowercase														=			config.to_json_string()

      def        A__ (		self		):







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







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







            """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=__lowerCAmelCase		)	, self.EXPECTED_OUTPUTS		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			BitsAndBytesConfig()
            lowercase														=			True

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

      def        A__ (		self		):







            """simple docstring"""



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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			BitsAndBytesConfig()

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

      def        A__ (		self		):







            """simple docstring"""



            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
            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        A__ (		self		):







            """simple docstring"""



            lowercase														=			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			):

      @classmethod
      def        A__ (		cls		):







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







            """simple docstring"""



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

      def        A__ (		self		):







            """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=__lowerCAmelCase	, device_map="""auto"""		)
            lowercase														=			self.tokenizer(self.input_text	, return_tensors="""pt"""		).to(0		)
            lowercase														=			model.generate(**__lowerCAmelCase		)

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







      def        A__ (		self		):







            """simple docstring"""



            import bitsandbytes as bnb

            from transformers import TaForConditionalGeneration

            # test with `t5-small`
            lowercase														=			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		)		)

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

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






class 				_A		( lowerCAmelCase			):

      def        A__ (		self		):







            """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=__lowerCAmelCase	, device_map="""auto"""		)
            # Sequence classification model
            lowercase														=			AutoModelForSequenceClassification.from_pretrained(
                self.model_name	, load_in_abit=__lowerCAmelCase	, device_map="""auto"""		)
            # CausalLM model
            lowercase														=			AutoModelForCausalLM.from_pretrained(self.model_name	, load_in_abit=__lowerCAmelCase	, device_map="""auto"""		)
            # Seq2seq model
            lowercase														=			AutoModelForSeqaSeqLM.from_pretrained(
                self.seq_to_seq_name	, load_in_abit=__lowerCAmelCase	, device_map="""auto"""		)

      def        A__ (		self		):







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







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







            """simple docstring"""



            super().setUp()

      def        A__ (		self		):







            """simple docstring"""



            del self.pipe

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







      def        A__ (		self		):







            """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 				_A		( lowerCAmelCase			):

      def        A__ (		self		):







            """simple docstring"""



            super().setUp()







      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			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
            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=__lowerCAmelCase		)	, self.EXPECTED_OUTPUTS		)






class 				_A		( lowerCAmelCase			):

      def        A__ (		self		):







            """simple docstring"""



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







      def        A__ (		self		):







            """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=__lowerCAmelCase		)

            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(__lowerCAmelCase		)		):
                        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(**__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			):
      snake_case__    :     Union[str, Any]													=			'gpt2-xl'
      snake_case__    :     List[str]													=			3.31_91_85_48_54_15_21_87




"""simple docstring"""






def     UpperCAmelCase__ (       lowerCAmelCase__  :int	)				->					bool:
      '''simple docstring'''


      lowercase														=			(1 + 2_4 * n) ** 0.5
      return ((1 + root) / 6) % 1 == 0


def     UpperCAmelCase__ (       lowerCAmelCase__  :int = 5_0_0_0	)				->					int:
      '''simple docstring'''


      lowercase														=			[(i * (3 * i - 1)) // 2 for i in range(1      ,	lowerCAmelCase__	)]
      for i, pentagonal_i in enumerate(lowerCAmelCase__	):
            for j in range(lowerCAmelCase__      ,	len(lowerCAmelCase__	)	):
                  lowercase														=			pentagonal_nums[j]
                  lowercase														=			pentagonal_i + pentagonal_j
                  lowercase														=			pentagonal_j - pentagonal_i
                  if is_pentagonal(lowerCAmelCase__	) and is_pentagonal(lowerCAmelCase__	):
                        return b

      return -1


if __name__ == "__main__":
       print(F"""{solution() = }""")



