





'''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 				snake_case_      (		SCREAMING_SNAKE_CASE__       ):







  '''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       __SCREAMING_SNAKE_CASE	(						nn.Module			):



      '''simple docstring'''






      def __init__(					self						,   lowerCamelCase						,   lowerCamelCase	):
        super().__init__()
        _snake_case     		=				module
        _snake_case     		=				nn.Sequential(
            nn.Linear(module.in_features						,   lowerCamelCase						,   bias=lowerCamelCase	)						,   nn.Linear(lowerCamelCase						,   module.out_features						,   bias=lowerCamelCase	)						,   )
        _snake_case     		=				(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 	UpperCamelCase(					self						,   lowerCamelCase						,   *lowerCamelCase						,   **lowerCamelCase	):
        return self.module(lowerCamelCase						,   *lowerCamelCase						,   **lowerCamelCase	) + self.adapter(lowerCamelCase	)


@require_bitsandbytes
@require_accelerate
@require_torch
@require_torch_gpu
@slow
class       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Tuple      =   '''bigscience/bloom-1b7'''

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

     UpperCAmelCase__       :						Dict      =   '''Hello my name is'''
     UpperCAmelCase__       :						Any      =   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__       :						Union[str, Any]      =   10

     def 	UpperCamelCase(					self	):
       # Models and tokenizer
       _snake_case     		=				AutoTokenizer.from_pretrained(self.model_name	)







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       super().setUp()

       # Models and tokenizer
       _snake_case     		=				AutoModelForCausalLM.from_pretrained(
           self.model_name						,   torch_dtype=torch.floataa						,   device_map="auto"	)
       _snake_case     		=				AutoModelForCausalLM.from_pretrained(self.model_name						,   load_in_abit=lowerCamelCase						,   device_map="auto"	)

     def 	UpperCamelCase(					self	):
       del self.model_fpaa
       del self.model_abit

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

     def 	UpperCamelCase(					self	):
       _snake_case     		=				self.model_abit.config

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

       _snake_case     		=				config.to_dict()
       _snake_case     		=				config.to_diff_dict()

       _snake_case     		=				config.to_json_string()

     def 	UpperCamelCase(					self	):
       from bitsandbytes.nn import Paramsabit

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

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

     def 	UpperCamelCase(					self	):
       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 	UpperCamelCase(					self	):
       _snake_case     		=				self.tokenizer(self.input_text						,   return_tensors="pt"	)
       _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=lowerCamelCase	)						,   self.EXPECTED_OUTPUTS	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				BitsAndBytesConfig()
       _snake_case     		=				True

       _snake_case     		=				AutoModelForCausalLM.from_pretrained(
           self.model_name						,   quantization_config=lowerCamelCase						,   device_map="auto"	)

       _snake_case     		=				self.tokenizer(self.input_text						,   return_tensors="pt"	)
       _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=lowerCamelCase	)						,   self.EXPECTED_OUTPUTS	)

     def 	UpperCamelCase(					self	):
       with self.assertRaises(lowerCamelCase	), tempfile.TemporaryDirectory() as tmpdirname:
         self.model_abit.save_pretrained(lowerCamelCase	)

     def 	UpperCamelCase(					self	):
       _snake_case     		=				BitsAndBytesConfig()

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

     def 	UpperCamelCase(					self	):
       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
       _snake_case     		=				self.tokenizer(self.input_text						,   return_tensors="pt"	)

       _snake_case     		=				self.model_fpaa.to(torch.floataa	)
       _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
       _snake_case     		=				self.model_fpaa.to("cpu"	)

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

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






     def 	UpperCamelCase(					self	):
       _snake_case     		=				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       __SCREAMING_SNAKE_CASE	(						unittest.TestCase			):



     '''simple docstring'''






     @classmethod
     def 	UpperCamelCase(					cls	):
       _snake_case     		=				"t5-small"
       _snake_case     		=				"google/flan-t5-small"  # flan-t5 uses dense-act instead of dense-relu-dense
       _snake_case     		=				AutoTokenizer.from_pretrained(cls.model_name	)
       _snake_case     		=				"Translate in German: Hello, my dog is cute"

     def 	UpperCamelCase(					self	):
       gc.collect()
       torch.cuda.empty_cache()

     def 	UpperCamelCase(					self	):
       from transformers import TaForConditionalGeneration

       _snake_case     		=				TaForConditionalGeneration._keep_in_fpaa_modules
       _snake_case     		=				None

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

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






     def 	UpperCamelCase(					self	):
       import bitsandbytes as bnb

       from transformers import TaForConditionalGeneration

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

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

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







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       super().setUp()
       # model_name
       _snake_case     		=				"bigscience/bloom-560m"
       _snake_case     		=				"t5-small"

       # Different types of model

       _snake_case     		=				AutoModel.from_pretrained(self.model_name						,   load_in_abit=lowerCamelCase						,   device_map="auto"	)
       # Sequence classification model
       _snake_case     		=				AutoModelForSequenceClassification.from_pretrained(
           self.model_name						,   load_in_abit=lowerCamelCase						,   device_map="auto"	)
       # CausalLM model
       _snake_case     		=				AutoModelForCausalLM.from_pretrained(self.model_name						,   load_in_abit=lowerCamelCase						,   device_map="auto"	)
       # Seq2seq model
       _snake_case     		=				AutoModelForSeqaSeqLM.from_pretrained(
           self.seq_to_seq_name						,   load_in_abit=lowerCamelCase						,   device_map="auto"	)

     def 	UpperCamelCase(					self	):
       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	):
       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       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       super().setUp()

     def 	UpperCamelCase(					self	):
       del self.pipe

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






     def 	UpperCamelCase(					self	):
       _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
       _snake_case     		=				self.pipe(self.input_text	)
       self.assertIn(pipeline_output[0]["generated_text"]						,   self.EXPECTED_OUTPUTS	)







@require_torch_multi_gpu
class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       super().setUp()






     def 	UpperCamelCase(					self	):
       _snake_case     		=				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
       _snake_case     		=				self.tokenizer(self.input_text						,   return_tensors="pt"	)

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







class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def 	UpperCamelCase(					self	):
       _snake_case     		=				"facebook/opt-350m"
       super().setUp()






     def 	UpperCamelCase(					self	):
       if version.parse(importlib.metadata.version("bitsandbytes"	)	) < version.parse("0.37.0"	):
         return

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

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

        # Step 3: dummy batch
       _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():
         _snake_case     		=				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       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     UpperCAmelCase__       :						Union[str, Any]      =   '''gpt2-xl'''
     UpperCAmelCase__       :						Dict      =   3.31_91_85_48_54_15_21_87









'''simple docstring'''



from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
 from PIL import Image

 from ..image_utils import load_image

if is_tf_available():
 from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
 import torch

 from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

__magic_name__   :						Optional[int]    =			logging.get_logger(__name__)


@add_end_docstrings(__UpperCamelCase			)
class       __SCREAMING_SNAKE_CASE	(						__UpperCamelCase			):



     '''simple docstring'''






     def __init__(					self						,   *lowerCamelCase						,   **lowerCamelCase	):
       super().__init__(*lowerCamelCase						,   **lowerCamelCase	)
       requires_backends(self						,   "vision"	)
       self.check_model_type(
           TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == "tf" else MODEL_FOR_VISION_2_SEQ_MAPPING	)

     def 	UpperCamelCase(					self						,   lowerCamelCase=None						,   lowerCamelCase=None						,   lowerCamelCase=None	):
       _snake_case     		=				{}
       _snake_case     		=				{}

       if prompt is not None:
         _snake_case     		=				prompt

       if generate_kwargs is not None:
         _snake_case     		=				generate_kwargs
       if max_new_tokens is not None:
         if "generate_kwargs" not in forward_kwargs:
           _snake_case     		=				{}
         if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
           raise ValueError(
               "'max_new_tokens' is defined twice, once in 'generate_kwargs' and once as a direct parameter,"
               " please use only one"	)
         _snake_case     		=				max_new_tokens
       return preprocess_params, forward_kwargs, {}

     def __call__(					self						,   lowerCamelCase						,   **lowerCamelCase	):
       return super().__call__(lowerCamelCase						,   **lowerCamelCase	)

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase=None	):
       _snake_case     		=				load_image(lowerCamelCase	)

       if prompt is not None:
         if not isinstance(lowerCamelCase						,   lowerCamelCase	):
           raise ValueError(
               F'''Received an invalid text input, got - {type(lowerCamelCase	)} - but expected a single string. '''
               "Note also that one single text can be provided for conditional image to text generation."	)

         _snake_case     		=				self.model.config.model_type

         if model_type == "git":
           _snake_case     		=				self.image_processor(images=lowerCamelCase						,   return_tensors=self.framework	)
           _snake_case     		=				self.tokenizer(text=lowerCamelCase						,   add_special_tokens=lowerCamelCase	).input_ids
           _snake_case     		=				[self.tokenizer.cls_token_id] + input_ids
           _snake_case     		=				torch.tensor(lowerCamelCase	).unsqueeze(0	)
           model_inputs.update({"input_ids": input_ids}	)

         elif model_type == "pix2struct":
           _snake_case     		=				self.image_processor(images=lowerCamelCase						,   header_text=lowerCamelCase						,   return_tensors=self.framework	)

         elif model_type != "vision-encoder-decoder":
           # vision-encoder-decoder does not support conditional generation
           _snake_case     		=				self.image_processor(images=lowerCamelCase						,   return_tensors=self.framework	)
           _snake_case     		=				self.tokenizer(lowerCamelCase						,   return_tensors=self.framework	)
           model_inputs.update(lowerCamelCase	)

         else:
           raise ValueError(F'''Model type {model_type} does not support conditional text generation'''	)

       else:
         _snake_case     		=				self.image_processor(images=lowerCamelCase						,   return_tensors=self.framework	)

       if self.model.config.model_type == "git" and prompt is None:
         _snake_case     		=				None

       return model_inputs

     def 	UpperCamelCase(					self						,   lowerCamelCase						,   lowerCamelCase=None	):
       # Git model sets `model_inputs["input_ids"] = None` in `preprocess` (when `prompt=None`). In batch model, the
       # pipeline will group them into a list of `None`, which fail `_forward`. Avoid this by checking it first.
       if (
           "input_ids" in model_inputs
           and isinstance(model_inputs["input_ids"]						,   lowerCamelCase	)
           and all(x is None for x in model_inputs["input_ids"]	)
       ):
         _snake_case     		=				None

       if generate_kwargs is None:
         _snake_case     		=				{}
       # FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
       #  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
       #  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
       #  in the `_prepare_model_inputs` method.
       _snake_case     		=				model_inputs.pop(self.model.main_input_name	)
       _snake_case     		=				self.model.generate(lowerCamelCase						,   **lowerCamelCase						,   **lowerCamelCase	)
       return model_outputs






     def 	UpperCamelCase(					self						,   lowerCamelCase	):
       _snake_case     		=				[]
       for output_ids in model_outputs:
         _snake_case     		=				{
             "generated_text": self.tokenizer.decode(
                 lowerCamelCase						,   skip_special_tokens=lowerCamelCase						,   )
         }
         records.append(lowerCamelCase	)
       return records


