





"""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  _SCREAMING_SNAKE_CASE						(      _lowercase		:       Optional[int]   )					->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 			__UpperCamelCase							(		nn.Module		):






									def __init__(  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					str:
										super().__init__()
										a :						Union[str, Any]      =							module
										a :						Any      =							nn.Sequential(
										    nn.Linear(module.in_features ,   lowerCAmelCase__ ,   bias=lowerCAmelCase__ ) ,   nn.Linear(lowerCAmelCase__ ,   module.out_features ,   bias=lowerCAmelCase__ ) ,   )
										a :						str      =							(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__ )      ->					int:
										return self.module(lowerCAmelCase__ ,   *lowerCAmelCase__ ,   **lowerCAmelCase__ ) + self.adapter(lowerCAmelCase__ )


@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class 			__UpperCamelCase							(		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
							lowerCamelCase  :   int    ="""bigscience/bloom-1b7"""

							# Constant values
							lowerCamelCase  :   Any    =2.109_659_552_692_574

							lowerCamelCase  :   Optional[int]    ="""Hello my name is"""
							lowerCamelCase  :   List[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"""		)
							lowerCamelCase  :   List[str]    =10






							def        __a  (  self )      ->					Any:
								# Models and tokenizer
								a :						Tuple      =							AutoTokenizer.from_pretrained(self.model_name )


class 			__UpperCamelCase							(		a__		):






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

								# Models and tokenizer
								a :						Tuple      =							AutoModelForCausalLM.from_pretrained(
								    self.model_name ,   torch_dtype=torch.floataa ,   device_map="auto" )
								a :						str      =							AutoModelForCausalLM.from_pretrained(self.model_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )






							def        __a  (  self )      ->					str:
								del self.model_fpaa
								del self.model_abit

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






							def        __a  (  self )      ->					Optional[Any]:
								a :						Union[str, Any]      =							self.model_abit.config

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

								a :						Tuple      =							config.to_dict()
								a :						List[Any]      =							config.to_diff_dict()

								a :						int      =							config.to_json_string()






							def        __a  (  self )      ->					Tuple:
								from bitsandbytes.nn import Paramsabit

								a :						List[str]      =							self.model_fpaa.get_memory_footprint()
								a :						int      =							self.model_abit.get_memory_footprint()

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






							def        __a  (  self )      ->					List[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(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 )      ->					Tuple:
								a :						str      =							self.tokenizer(self.input_text ,   return_tensors="pt" )
								a :						Optional[int]      =							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 )      ->					Union[str, Any]:
								a :						Dict      =							BitsAndBytesConfig()
								a :						List[str]      =							True

								a :						Tuple      =							AutoModelForCausalLM.from_pretrained(
								    self.model_name ,   quantization_config=lowerCAmelCase__ ,   device_map="auto" )

								a :						str      =							self.tokenizer(self.input_text ,   return_tensors="pt" )
								a :						Dict      =							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 )      ->					Optional[Any]:
								with self.assertRaises(lowerCAmelCase__ ), tempfile.TemporaryDirectory() as tmpdirname:
									self.model_abit.save_pretrained(lowerCAmelCase__ )






							def        __a  (  self )      ->					Dict:
								a :						Any      =							BitsAndBytesConfig()

								with self.assertRaises(lowerCAmelCase__ ):
									a :						int      =							AutoModelForCausalLM.from_pretrained(
									    self.model_name ,   quantization_config=lowerCAmelCase__ ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" ,   bnb_abit_quant_type="nf4" ,   )






							def        __a  (  self )      ->					List[Any]:
								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
								a :						int      =							self.tokenizer(self.input_text ,   return_tensors="pt" )

								a :						Union[str, Any]      =							self.model_fpaa.to(torch.floataa )
								a :						int      =							self.model_fpaa.generate(input_ids=encoded_input["input_ids"].to(0 ) ,   max_new_tokens=10 )

								# Check this does not throw an error
								a :						Optional[int]      =							self.model_fpaa.to("cpu" )

								# Check this does not throw an error
								a :						Union[str, Any]      =							self.model_fpaa.half()

								# Check this does not throw an error
								a :						List[str]      =							self.model_fpaa.float()






							def        __a  (  self )      ->					Dict:
								a :						Union[str, Any]      =							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 			__UpperCamelCase							(		unittest.TestCase		):






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






							def        __a  (  self )      ->					str:
								gc.collect()
								torch.cuda.empty_cache()






							def        __a  (  self )      ->					Union[str, Any]:
								from transformers import TaForConditionalGeneration

								a :						Dict      =							TaForConditionalGeneration._keep_in_fpaa_modules
								a :						Optional[Any]      =							None

								# test with `t5-small`
								a :						Optional[Any]      =							TaForConditionalGeneration.from_pretrained(self.model_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )
								a :						Dict      =							self.tokenizer(self.input_text ,   return_tensors="pt" ).to(0 )
								a :						Any      =							model.generate(**lowerCAmelCase__ )

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






							def        __a  (  self )      ->					Optional[Any]:
								import bitsandbytes as bnb

								from transformers import TaForConditionalGeneration

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

								a :						Any      =							self.tokenizer(self.input_text ,   return_tensors="pt" ).to(0 )
								a :						str      =							model.generate(**lowerCAmelCase__ )

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


class 			__UpperCamelCase							(		a__		):






							def        __a  (  self )      ->					str:
								super().setUp()
								# model_name
								a :						Union[str, Any]      =							"bigscience/bloom-560m"
								a :						Dict      =							"t5-small"

								# Different types of model

								a :						Optional[Any]      =							AutoModel.from_pretrained(self.model_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )
								# Sequence classification model
								a :						Tuple      =							AutoModelForSequenceClassification.from_pretrained(
								    self.model_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )
								# CausalLM model
								a :						List[str]      =							AutoModelForCausalLM.from_pretrained(self.model_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )
								# Seq2seq model
								a :						int      =							AutoModelForSeqaSeqLM.from_pretrained(
								    self.seq_to_seq_name ,   load_in_abit=lowerCAmelCase__ ,   device_map="auto" )






							def        __a  (  self )      ->					int:
								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 )      ->					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 			__UpperCamelCase							(		a__		):






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






							def        __a  (  self )      ->					int:
								del self.pipe

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






							def        __a  (  self )      ->					List[str]:
								a :						int      =							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
								a :						Dict      =							self.pipe(self.input_text )
								self.assertIn(pipeline_output[0]["generated_text"] ,   self.EXPECTED_OUTPUTS )


@require_torch_multi_gpu
class 			__UpperCamelCase							(		a__		):






							def        __a  (  self )      ->					Tuple:
								super().setUp()






							def        __a  (  self )      ->					int:
								a :						Union[str, Any]      =							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
								a :						Optional[Any]      =							self.tokenizer(self.input_text ,   return_tensors="pt" )

								# Second real batch
								a :						Optional[int]      =							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 			__UpperCamelCase							(		a__		):






							def        __a  (  self )      ->					Union[str, Any]:
								a :						List[Any]      =							"facebook/opt-350m"
								super().setUp()






							def        __a  (  self )      ->					Optional[int]:
								if version.parse(importlib.metadata.version("bitsandbytes" ) ) < version.parse("0.37.0" ):
									return

								# Step 1: freeze all parameters
								a :						Any      =							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():
									a :						Any      =							False  # freeze the model - train adapters later
									if param.ndim == 1:
										# cast the small parameters (e.g. layernorm) to fp32 for stability
										a :						List[Any]      =							param.data.to(torch.floataa )

        # Step 2: add adapters
								for _, module in model.named_modules():
									if "OPTAttention" in repr(type(lowerCAmelCase__ ) ):
										a :						Any      =							LoRALayer(module.q_proj ,   rank=16 )
										a :						List[Any]      =							LoRALayer(module.k_proj ,   rank=16 )
										a :						Union[str, Any]      =							LoRALayer(module.v_proj ,   rank=16 )

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

								# Step 4: Check if the gradient is not None
								with torch.cuda.amp.autocast():
									a :						str      =							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 			__UpperCamelCase							(		a__		):
							lowerCamelCase  :   Dict    ="""gpt2-xl"""
							lowerCamelCase  :   Tuple    =3.3_191_854_854_152_187









"""simple docstring"""




import unittest

from transformers import DonutProcessor


a		:					Optional[int]	      =     '''naver-clova-ix/donut-base'''


class 			__UpperCamelCase							(		unittest.TestCase		):






							def        __a  (  self )      ->					Dict:
								a :						Optional[Any]      =							DonutProcessor.from_pretrained(lowerCAmelCase__ )






							def        __a  (  self )      ->					List[Any]:
								a :						Optional[Any]      =							{
								    "name": "John Doe",
								    "age": "99",
								    "city": "Atlanta",
								    "state": "GA",
								    "zip": "30301",
								    "phone": "123-4567",
								    "nicknames": [{"nickname": "Johnny"}, {"nickname": "JD"}],
								}

								a :						Tuple      =							(
								    "<s_name>John Doe</s_name><s_age>99</s_age><s_city>Atlanta</s_city>"
								    "<s_state>GA</s_state><s_zip>30301</s_zip><s_phone>123-4567</s_phone>"
								    "<s_nicknames><s_nickname>Johnny</s_nickname>"
								    "<sep/><s_nickname>JD</s_nickname></s_nicknames>"
								)
								a :						List[Any]      =							self.processor.tokenajson(lowerCAmelCase__ )

								self.assertDictEqual(lowerCAmelCase__ ,   lowerCAmelCase__ )


