




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(							__magic_name__       )					->				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						:      Union[str, Any] ,		_A						:      nn.Module ,		_A						:      int )						->       Dict:






        """simple docstring"""







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






   def    __a     (    self						:      str ,		_A						:      Optional[Any] ,		*_A						:      str ,		**_A						:      List[str] )						->       Optional[int]:






        """simple docstring"""







        return self.module(_A ,		*_A ,		**_A ) + self.adapter(_A )
@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 		_A   (    unittest.TestCase ):
  _UpperCamelCase  :     List[str]					=     '''bigscience/bloom-1b7'''

  # Constant values
  _UpperCamelCase  :     List[Any]					=     2.1_09_65_95_52_69_25_74

  _UpperCamelCase  :     Dict					=     '''Hello my name is'''
  _UpperCamelCase  :     Dict					=     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  :     int					=     1_0






  def    __a     (    self						:      Dict )						->       Optional[Any]:






       """simple docstring"""







       lowercase   :  List[str]        =					AutoTokenizer.from_pretrained(self.model_name )


class 		_A   (    __A ):






  def    __a     (    self						:      Dict )						->       int:






       """simple docstring"""







       super().setUp()

       # Models and tokenizer
       lowercase   :  Dict        =					AutoModelForCausalLM.from_pretrained(
           self.model_name ,		torch_dtype=torch.floataa ,		device_map='''auto''' )
       lowercase   :  int        =					AutoModelForCausalLM.from_pretrained(self.model_name ,		load_in_abit=_A ,		device_map='''auto''' )






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






       """simple docstring"""







       del self.model_fpaa
       del self.model_abit

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






  def    __a     (    self						:      Tuple )						->       Tuple:






       """simple docstring"""







       lowercase   :  Optional[int]        =					self.model_abit.config

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

       lowercase   :  Union[str, Any]        =					config.to_dict()
       lowercase   :  str        =					config.to_diff_dict()

       lowercase   :  Union[str, Any]        =					config.to_json_string()






  def    __a     (    self						:      Dict )						->       Optional[int]:






       """simple docstring"""







       from bitsandbytes.nn import Paramsabit

       lowercase   :  int        =					self.model_fpaa.get_memory_footprint()
       lowercase   :  int        =					self.model_abit.get_memory_footprint()

       self.assertAlmostEqual(mem_fpaa / mem_abit ,		self.EXPECTED_RELATIVE_DIFFERENCE )
       lowercase   :  Union[str, Any]        =					get_some_linear_layer(self.model_abit )
       self.assertTrue(linear.weight.__class__ == Paramsabit )






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






       """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(_A ,		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						:      Dict )						->       Union[str, Any]:






       """simple docstring"""







       lowercase   :  Optional[int]        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' )
       lowercase   :  Tuple        =					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=_A ) ,		self.EXPECTED_OUTPUTS )






  def    __a     (    self						:      Optional[int] )						->       int:






       """simple docstring"""







       lowercase   :  Tuple        =					BitsAndBytesConfig()
       lowercase   :  Optional[int]        =					True

       lowercase   :  int        =					AutoModelForCausalLM.from_pretrained(
           self.model_name ,		quantization_config=_A ,		device_map='''auto''' )

       lowercase   :  int        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' )
       lowercase   :  Tuple        =					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=_A ) ,		self.EXPECTED_OUTPUTS )






  def    __a     (    self						:      Any )						->       str:






       """simple docstring"""







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






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






       """simple docstring"""







       lowercase   :  int        =					BitsAndBytesConfig()

       with self.assertRaises(_A ):
            lowercase   :  List[str]        =					AutoModelForCausalLM.from_pretrained(
                self.model_name ,		quantization_config=_A ,		load_in_abit=_A ,		device_map='''auto''' ,		bnb_abit_quant_type='''nf4''' ,		)






  def    __a     (    self						:      Any )						->       Tuple:






       """simple docstring"""







       with self.assertRaises(_A ):
            # Tries with `str`
            self.model_abit.to('''cpu''' )

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

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

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

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

       # Test if we did not break anything
       lowercase   :  List[str]        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' )

       lowercase   :  Optional[Any]        =					self.model_fpaa.to(torch.floataa )
       lowercase   :  Optional[Any]        =					self.model_fpaa.generate(input_ids=encoded_input['''input_ids'''].to(0 ) ,		max_new_tokens=10 )

       # Check this does not throw an error
       lowercase   :  Tuple        =					self.model_fpaa.to('''cpu''' )

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

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






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






       """simple docstring"""







       lowercase   :  int        =					AutoModelForSeqaSeqLM.from_pretrained('''t5-small''' ,		load_in_abit=_A ,		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						:      str )						->       str:






       """simple docstring"""







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






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






       """simple docstring"""







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






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






       """simple docstring"""







       from transformers import TaForConditionalGeneration

       lowercase   :  Any        =					TaForConditionalGeneration._keep_in_fpaa_modules
       lowercase   :  Dict        =					None

       # test with `t5-small`
       lowercase   :  List[Any]        =					TaForConditionalGeneration.from_pretrained(self.model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       lowercase   :  int        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' ).to(0 )
       lowercase   :  Any        =					model.generate(**_A )

       # test with `flan-t5-small`
       lowercase   :  Dict        =					TaForConditionalGeneration.from_pretrained(
           self.dense_act_model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       lowercase   :  int        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' ).to(0 )
       lowercase   :  Optional[int]        =					model.generate(**_A )
       lowercase   :  Tuple        =					modules






  def    __a     (    self						:      int )						->       Dict:






       """simple docstring"""







       import bitsandbytes as bnb

       from transformers import TaForConditionalGeneration

       # test with `t5-small`
       lowercase   :  Optional[Any]        =					TaForConditionalGeneration.from_pretrained(self.model_name ,		load_in_abit=_A ,		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   :  List[Any]        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' ).to(0 )
       lowercase   :  Optional[Any]        =					model.generate(**_A )

       # test with `flan-t5-small`
       lowercase   :  List[str]        =					TaForConditionalGeneration.from_pretrained(
           self.dense_act_model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       lowercase   :  Dict        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' ).to(0 )
       lowercase   :  Tuple        =					model.generate(**_A )


class 		_A   (    __A ):






  def    __a     (    self						:      Any )						->       Tuple:






       """simple docstring"""







       super().setUp()
       # model_name
       lowercase   :  Optional[int]        =					"bigscience/bloom-560m"
       lowercase   :  List[str]        =					"t5-small"

       # Different types of model

       lowercase   :  str        =					AutoModel.from_pretrained(self.model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       # Sequence classification model
       lowercase   :  Union[str, Any]        =					AutoModelForSequenceClassification.from_pretrained(
           self.model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       # CausalLM model
       lowercase   :  List[Any]        =					AutoModelForCausalLM.from_pretrained(self.model_name ,		load_in_abit=_A ,		device_map='''auto''' )
       # Seq2seq model
       lowercase   :  List[str]        =					AutoModelForSeqaSeqLM.from_pretrained(
           self.seq_to_seq_name ,		load_in_abit=_A ,		device_map='''auto''' )






  def    __a     (    self						:      Optional[int] )						->       Dict:






       """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						:      List[Any] )						->       List[str]:






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






  def    __a     (    self						:      Dict )						->       Dict:






       """simple docstring"""







       super().setUp()






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






       """simple docstring"""







       del self.pipe

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






  def    __a     (    self						:      Dict )						->       Dict:






       """simple docstring"""







       lowercase   :  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
       lowercase   :  Any        =					self.pipe(self.input_text )
       self.assertIn(pipeline_output[0]['''generated_text'''] ,		self.EXPECTED_OUTPUTS )


@require_torch_multi_gpu
class 		_A   (    __A ):






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






       """simple docstring"""







       super().setUp()






  def    __a     (    self						:      List[Any] )						->       Any:






       """simple docstring"""







       lowercase   :  List[Any]        =					AutoModelForCausalLM.from_pretrained(
           self.model_name ,		load_in_abit=_A ,		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   :  Tuple        =					self.tokenizer(self.input_text ,		return_tensors='''pt''' )

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


class 		_A   (    __A ):






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






       """simple docstring"""







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






  def    __a     (    self						:      int )						->       List[Any]:






       """simple docstring"""







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

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

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

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

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

        # Step 3: dummy batch
       lowercase   :  Dict        =					self.tokenizer('''Test batch ''' ,		return_tensors='''pt''' ).to(0 )

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

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


class 		_A   (    __A ):
  _UpperCamelCase  :     int					=     '''gpt2-xl'''
  _UpperCamelCase  :     Optional[Any]					=     3.31_91_85_48_54_15_21_87


import argparse
import logging
from collections import namedtuple

import torch
from model_bertabs import BertAbsSummarizer
from models.model_builder import AbsSummarizer  # The authors' implementation

from transformers import BertTokenizer


logging.basicConfig(level=logging.INFO)
a_      =     logging.getLogger(__name__)


a_      =     'Hello world! cécé herlolip'


a_      =     namedtuple(
    'BertAbsConfig',
    [
        'temp_dir',
        'large',
        'use_bert_emb',
        'finetune_bert',
        'encoder',
        'share_emb',
        'max_pos',
        'enc_layers',
        'enc_hidden_size',
        'enc_heads',
        'enc_ff_size',
        'enc_dropout',
        'dec_layers',
        'dec_hidden_size',
        'dec_heads',
        'dec_ff_size',
        'dec_dropout',
    ],
)

def 			lowerCamelCase__			(       _a				,      _a):
	SCREAMING_SNAKE_CASE :   List[Any]					  =							BertAbsConfig(
	    temp_dir="."				,      finetune_bert=_a				,      large=_a				,      share_emb=_a				,      use_bert_emb=_a				,      encoder="bert"				,      max_pos=512				,      enc_layers=6				,      enc_hidden_size=512				,      enc_heads=8				,      enc_ff_size=512				,      enc_dropout=0.2				,      dec_layers=6				,      dec_hidden_size=768				,      dec_heads=8				,      dec_ff_size=2048				,      dec_dropout=0.2				,      )
	SCREAMING_SNAKE_CASE :   Dict					  =							torch.load(_a				,      lambda _a				,      _a: storage)
	SCREAMING_SNAKE_CASE :   str					  =							AbsSummarizer(_a				,      torch.device("cpu")				,      _a)
	original.eval()

	SCREAMING_SNAKE_CASE :   List[str]					  =							BertAbsSummarizer(_a				,      torch.device("cpu"))
	new_model.eval()

	# -------------------
	# Convert the weights
	# -------------------

	logging.info("convert the model")
	new_model.bert.load_state_dict(original.bert.state_dict())
	new_model.decoder.load_state_dict(original.decoder.state_dict())
	new_model.generator.load_state_dict(original.generator.state_dict())

	# ----------------------------------
	# Make sure the outpus are identical
	# ----------------------------------

	logging.info("Make sure that the models' outputs are identical")
	SCREAMING_SNAKE_CASE :   List[str]					  =							BertTokenizer.from_pretrained("bert-base-uncased")

	# prepare the model inputs
	SCREAMING_SNAKE_CASE :   List[str]					  =							tokenizer.encode("This is sample éàalj'-.")
	encoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(_a)))
	SCREAMING_SNAKE_CASE :   int					  =							torch.tensor(_a).unsqueeze(0)
	SCREAMING_SNAKE_CASE :   List[Any]					  =							tokenizer.encode("This is sample 3 éàalj'-.")
	decoder_input_ids.extend([tokenizer.pad_token_id] * (512 - len(_a)))
	SCREAMING_SNAKE_CASE :   int					  =							torch.tensor(_a).unsqueeze(0)

	# failsafe to make sure the weights reset does not affect the
	# loaded weights.
	assert torch.max(torch.abs(original.generator[0].weight - new_model.generator[0].weight)) == 0

	# forward pass
	SCREAMING_SNAKE_CASE :   List[Any]					  =							encoder_input_ids
	SCREAMING_SNAKE_CASE :   List[Any]					  =							decoder_input_ids
	SCREAMING_SNAKE_CASE :   Dict					  =							None
	SCREAMING_SNAKE_CASE :   Optional[Any]					  =							None
	SCREAMING_SNAKE_CASE :   List[Any]					  =							None
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							None
	SCREAMING_SNAKE_CASE :   Dict					  =							None

	# The original model does not apply the geneator layer immediatly but rather in
	# the beam search (where it combines softmax + linear layer). Since we already
	# apply the softmax in our generation process we only apply the linear layer here.
	# We make sure that the outputs of the full stack are identical
	SCREAMING_SNAKE_CASE :   Optional[int]					  =							original(_a				,      _a				,      _a				,      _a				,      _a				,      _a				,      _a)[0]
	SCREAMING_SNAKE_CASE :   Dict					  =							original.generator(_a)

	SCREAMING_SNAKE_CASE :   Any					  =							new_model(
	    _a				,      _a				,      _a				,      _a				,      _a)[0]
	SCREAMING_SNAKE_CASE :   Tuple					  =							new_model.generator(_a)

	SCREAMING_SNAKE_CASE :   List[Any]					  =							torch.max(torch.abs(output_converted_model - output_original_model)).item()
	print("Maximum absolute difference beween weights: {:.2f}".format(_a))
	SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.max(torch.abs(output_converted_generator - output_original_generator)).item()
	print("Maximum absolute difference beween weights: {:.2f}".format(_a))

	SCREAMING_SNAKE_CASE :   int					  =							torch.allclose(_a				,      _a				,      atol=1E-3)
	if are_identical:
		logging.info("all weights are equal up to 1e-3")
	else:
		raise ValueError("the weights are different. The new model is likely different from the original one.")

	# The model has been saved with torch.save(model) and this is bound to the exact
	# directory structure. We save the state_dict instead.
	logging.info("saving the model's state dictionary")
	torch.save(
	    new_model.state_dict()				,      "./bertabs-finetuned-cnndm-extractive-abstractive-summarization/pytorch_model.bin")


if __name__ == "__main__":
			a_      =     argparse.ArgumentParser()
			parser.add_argument(
			    '--bertabs_checkpoint_path',
			    default=None,
			    type=str,
			    required=True,
			    help='Path the official PyTorch dump.',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path',
			    default=None,
			    type=str,
			    required=True,
			    help='Path to the output PyTorch model.',
			)
			a_      =     parser.parse_args()

			convert_bertabs_checkpoints(
			    args.bertabs_checkpoint_path,
			    args.pytorch_dump_folder_path,
			)