


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( a__):
				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]					,     _a       :   nn.Module					,     _a       :   int	)				->  List[Any]:


																	"""simple docstring"""





																	super().__init__()
																	_SCREAMING_SNAKE_CASE      	=module
																	_SCREAMING_SNAKE_CASE      	=nn.Sequential(
																	    nn.Linear(module.in_features					,     lowercase_					,     bias=lowercase_	)					,     nn.Linear(lowercase_					,     module.out_features					,     bias=lowercase_	)					,     )
																	_SCREAMING_SNAKE_CASE      	=(2.0 / (5 * min(module.in_features					,     module.out_features	))) ** 0.5
																	nn.init.normal_(self.adapter[0].weight					,     std=lowercase_	)
																	nn.init.zeros_(self.adapter[1].weight	)
																	self.adapter.to(module.weight.device	)




													def 				__UpperCamelCase					(							self       :   Any					,     _a       :   Optional[Any]					,     *_a       :   str					,     **_a       :   Union[str, Any]	)				->  Union[str, Any]:


																	"""simple docstring"""





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

@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class     A__				(		unittest.TestCase	):
						UpperCAmelCase   							=    "bigscience/bloom-1b7"

						# Constant values
						UpperCAmelCase   							=    2.1_09_65_95_52_69_25_74

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




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=AutoTokenizer.from_pretrained(self.model_name	)

class     A__				(		a_	):




						def 				__UpperCamelCase					(							self       :   str	)				->  List[str]:


										"""simple docstring"""





										super().setUp()

										# Models and tokenizer
										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(
										    self.model_name					,     torch_dtype=torch.floataa					,     device_map='''auto'''	)
										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(self.model_name					,     load_in_abit=lowercase_					,     device_map='''auto'''	)




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


										"""simple docstring"""





										del self.model_fpaa
										del self.model_abit

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




						def 				__UpperCamelCase					(							self       :   Tuple	)				->  int:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.model_abit.config

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

										_SCREAMING_SNAKE_CASE      	=config.to_dict()
										_SCREAMING_SNAKE_CASE      	=config.to_diff_dict()

										_SCREAMING_SNAKE_CASE      	=config.to_json_string()




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


										"""simple docstring"""





										from bitsandbytes.nn import Paramsabit

										_SCREAMING_SNAKE_CASE      	=self.model_fpaa.get_memory_footprint()
										_SCREAMING_SNAKE_CASE      	=self.model_abit.get_memory_footprint()

										self.assertAlmostEqual(mem_fpaa / mem_abit					,     self.EXPECTED_RELATIVE_DIFFERENCE	)
										_SCREAMING_SNAKE_CASE      	=get_some_linear_layer(self.model_abit	)
										self.assertTrue(linear.weight.__class__ == Paramsabit	)




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


										"""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(lowercase_					,     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       :   Optional[int]	)				->  Dict:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=self.tokenizer(self.input_text					,     return_tensors='''pt'''	)
										_SCREAMING_SNAKE_CASE      	=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=lowercase_	)					,     self.EXPECTED_OUTPUTS	)




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=BitsAndBytesConfig()
										_SCREAMING_SNAKE_CASE      	=True

										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(
										    self.model_name					,     quantization_config=lowercase_					,     device_map='''auto'''	)

										_SCREAMING_SNAKE_CASE      	=self.tokenizer(self.input_text					,     return_tensors='''pt'''	)
										_SCREAMING_SNAKE_CASE      	=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=lowercase_	)					,     self.EXPECTED_OUTPUTS	)




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


										"""simple docstring"""





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




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=BitsAndBytesConfig()

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




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


										"""simple docstring"""





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

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

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

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

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

										# Test if we did not break anything
										_SCREAMING_SNAKE_CASE      	=self.tokenizer(self.input_text					,     return_tensors='''pt'''	)

										_SCREAMING_SNAKE_CASE      	=self.model_fpaa.to(torch.floataa	)
										_SCREAMING_SNAKE_CASE      	=self.model_fpaa.generate(input_ids=encoded_input['''input_ids'''].to(0	)					,     max_new_tokens=10	)

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

										# Check this does not throw an error
										_SCREAMING_SNAKE_CASE      	=self.model_fpaa.half()

										# Check this does not throw an error
										_SCREAMING_SNAKE_CASE      	=self.model_fpaa.float()




						def 				__UpperCamelCase					(							self       :   List[str]	)				->  str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=AutoModelForSeqaSeqLM.from_pretrained('''t5-small'''					,     load_in_abit=lowercase_					,     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 				__UpperCamelCase					(							cls       :   List[Any]	)				->  Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	='''t5-small'''
										_SCREAMING_SNAKE_CASE      	='''google/flan-t5-small'''  # flan-t5 uses dense-act instead of dense-relu-dense
										_SCREAMING_SNAKE_CASE      	=AutoTokenizer.from_pretrained(cls.model_name	)
										_SCREAMING_SNAKE_CASE      	='''Translate in German: Hello, my dog is cute'''




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


										"""simple docstring"""





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




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


										"""simple docstring"""





										from transformers import TaForConditionalGeneration

										_SCREAMING_SNAKE_CASE      	=TaForConditionalGeneration._keep_in_fpaa_modules
										_SCREAMING_SNAKE_CASE      	=None

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

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




						def 				__UpperCamelCase					(							self       :   Optional[int]	)				->  List[str]:


										"""simple docstring"""





										import bitsandbytes as bnb

										from transformers import TaForConditionalGeneration

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

										_SCREAMING_SNAKE_CASE      	=self.tokenizer(self.input_text					,     return_tensors='''pt'''	).to(0	)
										_SCREAMING_SNAKE_CASE      	=model.generate(**lowercase_	)

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

class     A__				(		a_	):




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


										"""simple docstring"""





										super().setUp()
										# model_name
										_SCREAMING_SNAKE_CASE      	='''bigscience/bloom-560m'''
										_SCREAMING_SNAKE_CASE      	='''t5-small'''

										# Different types of model

										_SCREAMING_SNAKE_CASE      	=AutoModel.from_pretrained(self.model_name					,     load_in_abit=lowercase_					,     device_map='''auto'''	)
										# Sequence classification model
										_SCREAMING_SNAKE_CASE      	=AutoModelForSequenceClassification.from_pretrained(
										    self.model_name					,     load_in_abit=lowercase_					,     device_map='''auto'''	)
										# CausalLM model
										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(self.model_name					,     load_in_abit=lowercase_					,     device_map='''auto'''	)
										# Seq2seq model
										_SCREAMING_SNAKE_CASE      	=AutoModelForSeqaSeqLM.from_pretrained(
										    self.seq_to_seq_name					,     load_in_abit=lowercase_					,     device_map='''auto'''	)




						def 				__UpperCamelCase					(							self       :   List[Any]	)				->  Tuple:


										"""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 				__UpperCamelCase					(							self       :   Optional[int]	)				->  Optional[Any]:


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


										"""simple docstring"""





										super().setUp()




						def 				__UpperCamelCase					(							self       :   Tuple	)				->  Tuple:


										"""simple docstring"""





										del self.pipe

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




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


										"""simple docstring"""





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

@require_torch_multi_gpu
class     A__				(		a_	):




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


										"""simple docstring"""





										super().setUp()




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(
										    self.model_name					,     load_in_abit=lowercase_					,     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
										_SCREAMING_SNAKE_CASE      	=self.tokenizer(self.input_text					,     return_tensors='''pt'''	)

										# Second real batch
										_SCREAMING_SNAKE_CASE      	=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=lowercase_	)					,     self.EXPECTED_OUTPUTS	)

class     A__				(		a_	):




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	='''facebook/opt-350m'''
										super().setUp()




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


										"""simple docstring"""





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

										# Step 1: freeze all parameters
										_SCREAMING_SNAKE_CASE      	=AutoModelForCausalLM.from_pretrained(self.model_name					,     load_in_abit=lowercase_	)

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

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

        # Step 2: add adapters
										for _, module in model.named_modules():
														if "OPTAttention" in repr(type(lowercase_	)	):
																		_SCREAMING_SNAKE_CASE      	=LoRALayer(module.q_proj					,     rank=16	)
																		_SCREAMING_SNAKE_CASE      	=LoRALayer(module.k_proj					,     rank=16	)
																		_SCREAMING_SNAKE_CASE      	=LoRALayer(module.v_proj					,     rank=16	)

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

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

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

class     A__				(		a_	):
						UpperCAmelCase   							=    "gpt2-xl"
						UpperCAmelCase   							=    3.31_91_85_48_54_15_21_87


"""simple docstring"""





# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ..models.auto import AutoModelForVisionaSeq
from ..utils import requires_backends
from .base import PipelineTool


if TYPE_CHECKING:
							from PIL import Image







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			"Salesforce/blip-image-captioning-base"
					lowercase__												=			(
					    "This is a tool that generates a description of an image. It takes an input named `image` which should be the "
					    "image to caption, and returns a text that contains the description in English."
					)
					lowercase__												=			"image_captioner"
					lowercase__												=			AutoModelForVisionaSeq

					lowercase__												=			["image"]
					lowercase__												=			["text"]







					def __init__(			self     :							int  ,*lowercase_     :							int  ,**lowercase_     :							List[str]						):
						requires_backends(self  ,['''vision''']						)
						super().__init__(*lowercase_  ,**lowercase_						)







					def 			__lowerCAmelCase					(			self     :							List[str]  ,lowercase_     :							"Image"						):
						return self.pre_processor(images=lowercase_  ,return_tensors='''pt'''						)







					def 			__lowerCAmelCase					(			self     :							List[Any]  ,lowercase_     :							Any						):
						return self.model.generate(**lowercase_						)







					def 			__lowerCAmelCase					(			self     :							Dict  ,lowercase_     :							Union[str, Any]						):
						return self.pre_processor.batch_decode(lowercase_  ,skip_special_tokens=lowercase_						)[0].strip()




