
import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       GPTNeoXForCausalLM,
       GPTNeoXForQuestionAnswering,
       GPTNeoXForSequenceClassification,
       GPTNeoXForTokenClassification,
       GPTNeoXModel,
   )







class 					lowerCamelCase		:





     def __init__(self				:   int    ,   _A				:   Tuple    ,   _A				:   Union[str, Any]=1_3    ,   _A				:   Optional[Any]=7    ,   _A				:   str=True    ,   _A				:   Optional[int]=True    ,   _A				:   str=True    ,   _A				:   int=True    ,   _A				:   str=9_9    ,   _A				:   List[Any]=6_4    ,   _A				:   Tuple=5    ,   _A				:   Any=4    ,   _A				:   int=3_7    ,   _A				:   List[str]="gelu"    ,   _A				:   Optional[Any]=0.1    ,   _A				:   Optional[Any]=0.1    ,   _A				:   Optional[int]=5_1_2    ,   _A				:   Tuple=1_6    ,   _A				:   Union[str, Any]=2    ,   _A				:   Optional[Any]=0.02    ,   _A				:   List[str]=3    ,   _A				:   Any=4    ,   _A				:   List[Any]=None    ,   )						->						Optional[Any]:
         snake_case								= parent
         snake_case								= batch_size
         snake_case								= seq_length
         snake_case								= is_training
         snake_case								= use_input_mask
         snake_case								= use_token_type_ids
         snake_case								= use_labels
         snake_case								= vocab_size
         snake_case								= hidden_size
         snake_case								= num_hidden_layers
         snake_case								= num_attention_heads
         snake_case								= intermediate_size
         snake_case								= hidden_act
         snake_case								= hidden_dropout_prob
         snake_case								= attention_probs_dropout_prob
         snake_case								= max_position_embeddings
         snake_case								= type_vocab_size
         snake_case								= type_sequence_label_size
         snake_case								= initializer_range
         snake_case								= num_labels
         snake_case								= num_choices
         snake_case								= scope
         snake_case								= vocab_size - 1





     def 		UpperCAmelCase(self				:   Tuple       )						->						Any:
         snake_case								= ids_tensor([self.batch_size, self.seq_length]    ,   self.vocab_size       )

         snake_case								= None
         if self.use_input_mask:
             snake_case								= random_attention_mask([self.batch_size, self.seq_length]       )

         snake_case								= None
         if self.use_labels:
             snake_case								= ids_tensor([self.batch_size, self.seq_length]    ,   self.num_labels       )

         snake_case								= self.get_config()

         return config, input_ids, input_mask, token_labels





     def 		UpperCAmelCase(self				:   Any       )						->						Any:
         return GPTNeoXConfig(
             vocab_size=self.vocab_size    ,   hidden_size=self.hidden_size    ,   num_hidden_layers=self.num_hidden_layers    ,   num_attention_heads=self.num_attention_heads    ,   intermediate_size=self.intermediate_size    ,   hidden_act=self.hidden_act    ,   hidden_dropout_prob=self.hidden_dropout_prob    ,   attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,   max_position_embeddings=self.max_position_embeddings    ,   type_vocab_size=self.type_vocab_size    ,   is_decoder=_A    ,   initializer_range=self.initializer_range    ,   pad_token_id=self.pad_token_id    ,   )





     def 		UpperCAmelCase(self				:   List[Any]       )						->						List[str]:
         snake_case					,							snake_case					,							snake_case					,							snake_case								= self.prepare_config_and_inputs()

         snake_case								= True

         return config, input_ids, input_mask, token_labels





     def 		UpperCAmelCase(self				:   Optional[int]    ,   _A				:   int    ,   _A				:   Tuple    ,   _A				:   Tuple       )						->						Any:
         snake_case								= GPTNeoXModel(config=_A       )
         model.to(_A       )
         model.eval()
         snake_case								= model(_A    ,   attention_mask=_A       )
         snake_case								= model(_A       )
         self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, self.seq_length, self.hidden_size)       )





     def 		UpperCAmelCase(self				:   Optional[Any]    ,   _A				:   int    ,   _A				:   int    ,   _A				:   Optional[Any]       )						->						str:
         snake_case								= True
         snake_case								= GPTNeoXModel(_A       )
         model.to(_A       )
         model.eval()
         snake_case								= model(_A    ,   attention_mask=_A       )
         self.parent.assertEqual(result.last_hidden_state.shape    ,   (self.batch_size, self.seq_length, self.hidden_size)       )





     def 		UpperCAmelCase(self				:   str    ,   _A				:   Optional[int]    ,   _A				:   List[Any]    ,   _A				:   str    ,   _A				:   Union[str, Any]       )						->						List[str]:
         snake_case								= GPTNeoXForCausalLM(config=_A       )
         model.to(_A       )
         model.eval()
         snake_case								= model(_A    ,   attention_mask=_A    ,   labels=_A       )
         self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.seq_length, self.vocab_size)       )





     def 		UpperCAmelCase(self				:   Dict    ,   _A				:   Union[str, Any]    ,   _A				:   List[str]    ,   _A				:   int    ,   _A				:   List[str]       )						->						Tuple:
         snake_case								= self.num_labels
         snake_case								= GPTNeoXForQuestionAnswering(_A       )
         model.to(_A       )
         model.eval()
         snake_case								= model(_A    ,   attention_mask=_A       )
         self.parent.assertEqual(result.start_logits.shape    ,   (self.batch_size, self.seq_length)       )
         self.parent.assertEqual(result.end_logits.shape    ,   (self.batch_size, self.seq_length)       )





     def 		UpperCAmelCase(self				:   int    ,   _A				:   Any    ,   _A				:   Tuple    ,   _A				:   Any    ,   _A				:   Dict       )						->						Dict:
         snake_case								= self.num_labels
         snake_case								= GPTNeoXForSequenceClassification(_A       )
         model.to(_A       )
         model.eval()
         snake_case								= ids_tensor([self.batch_size]    ,   self.type_sequence_label_size       )
         snake_case								= model(_A    ,   attention_mask=_A    ,   labels=_A       )
         self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.num_labels)       )





     def 		UpperCAmelCase(self				:   Any    ,   _A				:   Tuple    ,   _A				:   Union[str, Any]    ,   _A				:   Optional[int]    ,   _A				:   Dict       )						->						Any:
         snake_case								= self.num_labels
         snake_case								= GPTNeoXForTokenClassification(_A       )
         model.to(_A       )
         model.eval()
         snake_case								= model(_A    ,   attention_mask=_A    ,   labels=_A       )
         self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.seq_length, self.num_labels)       )





     def 		UpperCAmelCase(self				:   Tuple    ,   _A				:   Optional[int]    ,   _A				:   Optional[Any]    ,   _A				:   Optional[Any]       )						->						int:
         snake_case								= True
         snake_case								= GPTNeoXForCausalLM(config=_A       )
         model.to(_A       )
         model.eval()

         # first forward pass
         snake_case								= model(_A    ,   attention_mask=_A    ,   use_cache=_A       )
         snake_case								= outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         snake_case								= ids_tensor((self.batch_size, 3)    ,   config.vocab_size       )
         snake_case								= ids_tensor((self.batch_size, 3)    ,   vocab_size=2       )

         # append to next input_ids and
         snake_case								= torch.cat([input_ids, next_tokens]    ,   dim=-1       )
         snake_case								= torch.cat([input_mask, next_mask]    ,   dim=-1       )

         snake_case								= model(_A    ,   attention_mask=_A    ,   output_hidden_states=_A       )
         snake_case								= output_from_no_past["hidden_states"][0]
         snake_case								= model(
             _A    ,   attention_mask=_A    ,   past_key_values=_A    ,   output_hidden_states=_A    ,   )["hidden_states"][0]

         # select random slice
         snake_case								= ids_tensor((1,)    ,   output_from_past.shape[-1]       ).item()
         snake_case								= output_from_no_past[:, -3:, random_slice_idx].detach()
         snake_case								= output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]       )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(_A    ,   _A    ,   atol=1E-3       )       )





     def 		UpperCAmelCase(self				:   Optional[Any]       )						->						Tuple:
         snake_case								= self.prepare_config_and_inputs()
         snake_case					,							snake_case					,							snake_case					,							snake_case								= config_and_inputs
         snake_case								= {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict







@require_torch
class 					lowerCamelCase		(  A_							,			A_							,			A_							,			unittest.TestCase  ):
     UpperCAmelCase__    :							List[str]									=		(
         (
             GPTNeoXModel,
             GPTNeoXForCausalLM,
             GPTNeoXForQuestionAnswering,
             GPTNeoXForSequenceClassification,
             GPTNeoXForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     UpperCAmelCase__    :							Dict									=		(GPTNeoXForCausalLM,) if is_torch_available() else ()
     UpperCAmelCase__    :							str									=		(
         {
             "feature-extraction": GPTNeoXModel,
             "question-answering": GPTNeoXForQuestionAnswering,
             "text-classification": GPTNeoXForSequenceClassification,
             "text-generation": GPTNeoXForCausalLM,
             "token-classification": GPTNeoXForTokenClassification,
             "zero-shot": GPTNeoXForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     UpperCAmelCase__    :							Optional[int]									=		False
     UpperCAmelCase__    :							List[str]									=		False
     UpperCAmelCase__    :							Dict									=		False
     UpperCAmelCase__    :							List[Any]									=		False





     def 		UpperCAmelCase(self				:   str       )						->						Optional[int]:
         snake_case								= GPTNeoXModelTester(self       )
         snake_case								= ConfigTester(self    ,   config_class=_A    ,   hidden_size=6_4    ,   num_attention_heads=8       )





     def 		UpperCAmelCase(self				:   Optional[int]       )						->						int:
         self.config_tester.run_common_tests()





     def 		UpperCAmelCase(self				:   List[str]       )						->						str:
         snake_case					,							snake_case					,							snake_case					,							snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(_A    ,   _A    ,   _A       )





     def 		UpperCAmelCase(self				:   Union[str, Any]       )						->						Any:
         snake_case					,							snake_case					,							snake_case					,							snake_case								= self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(_A    ,   _A    ,   _A       )





     def 		UpperCAmelCase(self				:   Optional[Any]       )						->						Dict:
         # This regression test was failing with PyTorch < 1.3
         snake_case					,							snake_case					,							snake_case					,							snake_case								= self.model_tester.prepare_config_and_inputs_for_decoder()

         snake_case								= None

         self.model_tester.create_and_check_model_as_decoder(_A    ,   _A    ,   _A       )





     def 		UpperCAmelCase(self				:   List[str]       )						->						Dict:
         snake_case					,							snake_case					,							snake_case					,							snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(_A    ,   _A    ,   _A       )





     def 		UpperCAmelCase(self				:   Dict       )						->						Any:
         snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_causal_lm(*_A       )





     def 		UpperCAmelCase(self				:   Tuple       )						->						Any:
         snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*_A       )





     def 		UpperCAmelCase(self				:   Dict       )						->						str:
         snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*_A       )





     def 		UpperCAmelCase(self				:   str       )						->						Any:
         snake_case								= self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*_A       )





     @unittest.skip(reason="Feed forward chunking is not implemented"       )
     def 		UpperCAmelCase(self				:   List[Any]       )						->						List[str]:
         pass





     @parameterized.expand([("linear",), ("dynamic",)]       )
     def 		UpperCAmelCase(self				:   Optional[int]    ,   _A				:   List[str]       )						->						int:
         snake_case					,							snake_case								= self.model_tester.prepare_config_and_inputs_for_common()
         snake_case								= ids_tensor([1, 1_0]    ,   config.vocab_size       )
         snake_case								= ids_tensor([1, int(config.max_position_embeddings * 1.5       )]    ,   config.vocab_size       )

         set_seed(4_2       )  # Fixed seed at init time so the two models get the same random weights
         snake_case								= GPTNeoXModel(_A       )
         original_model.to(_A       )
         original_model.eval()
         snake_case								= original_model(_A       ).last_hidden_state
         snake_case								= original_model(_A       ).last_hidden_state

         set_seed(4_2       )  # Fixed seed at init time so the two models get the same random weights
         snake_case								= {"type": scaling_type, "factor": 10.0}
         snake_case								= GPTNeoXModel(_A       )
         scaled_model.to(_A       )
         scaled_model.eval()
         snake_case								= scaled_model(_A       ).last_hidden_state
         snake_case								= scaled_model(_A       ).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
             self.assertTrue(torch.allclose(_A    ,   _A    ,   atol=1E-5       )       )
         else:
             self.assertFalse(torch.allclose(_A    ,   _A    ,   atol=1E-5       )       )

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(_A    ,   _A    ,   atol=1E-5       )       )







@require_torch
class 					lowerCamelCase		(  unittest.TestCase  ):





     @slow
     def 		UpperCAmelCase(self				:   Optional[int]       )						->						Optional[int]:
         snake_case								= AutoTokenizer.from_pretrained("EleutherAI/pythia-410m-deduped"       )
         for checkpointing in [True, False]:
             snake_case								= GPTNeoXForCausalLM.from_pretrained("EleutherAI/pythia-410m-deduped"       )

             if checkpointing:
                 model.gradient_checkpointing_enable()
             else:
                 model.gradient_checkpointing_disable()
             model.to(_A       )

             snake_case								= tokenizer("My favorite food is"    ,   return_tensors="pt"       ).to(_A       )
             # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
             # See: https://github.com/huggingface/transformers/pull/24193
             snake_case								= "My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"

             snake_case								= model.generate(**_A    ,   do_sample=_A    ,   max_new_tokens=2_0       )
             snake_case								= tokenizer.batch_decode(_A       )[0]

             self.assertEqual(_A    ,   _A       )



import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
   import bitsandbytes as bnb
   import torch
   import torch.nn as nn

   from ..pytorch_utils import ConvaD

if is_accelerate_available():
   from accelerate import init_empty_weights
   from accelerate.utils import find_tied_parameters

_A   =			logging.get_logger(__name__)



def 						lowercase_      (       A__			,		A__			,		A__			,		A__=None			,		A__=None		)							->							str:
    """simple docstring"""

    if "." in tensor_name:
        snake_case								= tensor_name.split("."		)
        for split in splits[:-1]:
            snake_case								= getattr(A__			,		A__		)
            if new_module is None:
                raise ValueError(F'{module} has no attribute {split}.'		)
            snake_case								= new_module
        snake_case								= splits[-1]

    if tensor_name not in module._parameters and tensor_name not in module._buffers:
        raise ValueError(F'{module} does not have a parameter or a buffer named {tensor_name}.'		)
    snake_case								= tensor_name in module._buffers
    snake_case								= getattr(A__			,		A__		)

    if old_value.device == torch.device("meta"		) and device not in ["meta", torch.device("meta"		)] and value is None:
        raise ValueError(F'{tensor_name} is on the meta device, we need a `value` to put in on {device}.'		)

    snake_case								= False
    snake_case								= False
    if is_buffer or not is_bitsandbytes_available():
        snake_case								= False
        snake_case								= False
    else:
        snake_case								= hasattr(bnb.nn			,		"Params4bit"		) and isinstance(module._parameters[tensor_name]			,		bnb.nn.Paramsabit		)
        snake_case								= isinstance(module._parameters[tensor_name]			,		bnb.nn.IntaParams		)

    if is_abit or is_abit:
        snake_case								= module._parameters[tensor_name]
        if param.device.type != "cuda":
            if value is None:
                snake_case								= old_value.to(A__		)
            elif isinstance(A__			,		torch.Tensor		):
                snake_case								= value.to("cpu"		)
                if value.dtype == torch.inta:
                    snake_case								= version.parse(importlib.metadata.version("bitsandbytes"		)		) > version.parse(
                        "0.37.2"		)
                    if not is_abit_serializable:
                        raise ValueError(
                            "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
                            "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."		)
            else:
                snake_case								= torch.tensor(A__			,		device="cpu"		)

            # Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
            # Since weights are saved in the correct "orientation", we skip transposing when loading.
            if issubclass(module.source_cls			,		A__		) and fpaa_statistics is None:
                snake_case								= new_value.T

            snake_case								= old_value.__dict__
            if is_abit:
                snake_case								= bnb.nn.IntaParams(A__			,		requires_grad=A__			,		**A__		).to(A__		)
            elif is_abit:
                snake_case								= bnb.nn.Paramsabit(A__			,		requires_grad=A__			,		**A__		).to(A__		)

            snake_case								= new_value
            if fpaa_statistics is not None:
                setattr(module.weight			,		"SCB"			,		fpaa_statistics.to(A__		)		)

    else:
        if value is None:
            snake_case								= old_value.to(A__		)
        elif isinstance(A__			,		torch.Tensor		):
            snake_case								= value.to(A__		)
        else:
            snake_case								= torch.tensor(A__			,		device=A__		)

        if is_buffer:
            snake_case								= new_value
        else:
            snake_case								= nn.Parameter(A__			,		requires_grad=old_value.requires_grad		)
            snake_case								= new_value



def 						lowercase_      (       A__			,		A__=None			,		A__=None			,		A__=None			,		A__=False		)							->							Optional[Any]:
    """simple docstring"""

    for name, module in model.named_children():
        if current_key_name is None:
            snake_case								= []
        current_key_name.append(A__		)

        if (isinstance(A__			,		nn.Linear		) or isinstance(A__			,		A__		)) and name not in modules_to_not_convert:
            # Check if the current key is not in the `modules_to_not_convert`
            if not any(key in ".".join(A__		) for key in modules_to_not_convert		):
                with init_empty_weights():
                    if isinstance(A__			,		A__		):
                        snake_case					,							snake_case								= module.weight.shape
                    else:
                        snake_case								= module.in_features
                        snake_case								= module.out_features

                    if quantization_config.quantization_method() == "llm_int8":
                        snake_case								= bnb.nn.LinearabitLt(
                            A__			,		A__			,		module.bias is not None			,		has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight			,		threshold=quantization_config.llm_inta_threshold			,		)
                        snake_case								= True
                    else:
                        if (
                            quantization_config.llm_inta_skip_modules is not None
                            and name in quantization_config.llm_inta_skip_modules
                        ):
                            pass
                        else:
                            snake_case								= bnb.nn.Linearabit(
                                A__			,		A__			,		module.bias is not None			,		quantization_config.bnb_abit_compute_dtype			,		compress_statistics=quantization_config.bnb_abit_use_double_quant			,		quant_type=quantization_config.bnb_abit_quant_type			,		)
                            snake_case								= True
                    # Store the module class in case we need to transpose the weight later
                    snake_case								= type(A__		)
                    # Force requires grad to False to avoid unexpected errors
                    model._modules[name].requires_grad_(A__		)
        if len(list(module.children()		)		) > 0:
            snake_case					,							snake_case								= _replace_with_bnb_linear(
                A__			,		A__			,		A__			,		A__			,		has_been_replaced=A__			,		)
        # Remove the last key for recursion
        current_key_name.pop(-1		)
    return model, has_been_replaced



def 						lowercase_      (       A__			,		A__=None			,		A__=None			,		A__=None		)							->							List[str]:
    """simple docstring"""

    snake_case								= ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
    snake_case					,							snake_case								= _replace_with_bnb_linear(
        A__			,		A__			,		A__			,		A__		)

    if not has_been_replaced:
        logger.warning(
            "You are loading your model in 8bit or 4bit but no linear modules were found in your model."
            " Please double check your model architecture, or submit an issue on github if you think this is"
            " a bug."		)

    return model



def 						lowercase_      (       *A__			,		**A__		)							->							Optional[int]:
    """simple docstring"""

    warnings.warn(
        "`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead"			,		A__			,		)
    return replace_with_bnb_linear(*A__			,		**A__		)



def 						lowercase_      (       *A__			,		**A__		)							->							Any:
    """simple docstring"""

    warnings.warn(
        "`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead"			,		A__			,		)
    return set_module_quantized_tensor_to_device(*A__			,		**A__		)





def 						lowercase_      (       A__		)							->							Union[str, Any]:
    """simple docstring"""

    snake_case								= deepcopy(A__		)  # this has 0 cost since it is done inside `init_empty_weights` context manager`
    tied_model.tie_weights()

    snake_case								= find_tied_parameters(A__		)
    # For compatibility with Accelerate < 0.18
    if isinstance(A__			,		A__		):
        snake_case								= sum(list(tied_params.values()		)			,		[]		) + list(tied_params.keys()		)
    else:
        snake_case								= sum(A__			,		[]		)
    snake_case								= len(A__		) > 0

    # Check if it is a base model
    snake_case								= not hasattr(A__			,		model.base_model_prefix		)

    # Ignore this for base models (BertModel, GPT2Model, etc.)
    if (not has_tied_params) and is_base_model:
        return []

    # otherwise they have an attached head
    snake_case								= list(model.named_children()		)
    snake_case								= [list_modules[-1][0]]

    # add last module together with tied weights
    snake_case								= set(A__		) - set(A__		)
    snake_case								= list(set(A__		)		) + list(A__		)

    # remove ".weight" from the keys
    snake_case								= [".weight", ".bias"]
    snake_case								= []
    for name in list_untouched:
        for name_to_remove in names_to_remove:
            if name_to_remove in name:
                snake_case								= name.replace(A__			,		""		)
        filtered_module_names.append(A__		)

    return filtered_module_names

