


"""simple docstring"""







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__    :							Dict            =					logging.get_logger(__name__)




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None       ):




					'''simple docstring'''


					if "." in tensor_name:
										__SCREAMING_SNAKE_CASE						       =    tensor_name.split("."       )
										for split in splits[:-1]:
															__SCREAMING_SNAKE_CASE						       =    getattr(lowerCAmelCase_  ,						lowerCAmelCase_       )
															if new_module is None:
																				raise ValueError(f"""{module} has no attribute {split}."""       )
															__SCREAMING_SNAKE_CASE						       =    new_module
										__SCREAMING_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}."""       )
					__SCREAMING_SNAKE_CASE						       =    tensor_name in module._buffers
					__SCREAMING_SNAKE_CASE						       =    getattr(lowerCAmelCase_  ,						lowerCAmelCase_       )

					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}."""       )

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

					if is_abit or is_abit:
										__SCREAMING_SNAKE_CASE						       =    module._parameters[tensor_name]
										if param.device.type != "cuda":
															if value is None:
																				__SCREAMING_SNAKE_CASE						       =    old_value.to(lowerCAmelCase_       )
															elif isinstance(lowerCAmelCase_  ,						torch.Tensor       ):
																				__SCREAMING_SNAKE_CASE						       =    value.to("cpu"       )
																				if value.dtype == torch.inta:
																									__SCREAMING_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:
																				__SCREAMING_SNAKE_CASE						       =    torch.tensor(lowerCAmelCase_  ,						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  ,						lowerCAmelCase_       ) and fpaa_statistics is None:
																				__SCREAMING_SNAKE_CASE						       =    new_value.T

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

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

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

										if is_buffer:
															__SCREAMING_SNAKE_CASE						       =    new_value
										else:
															__SCREAMING_SNAKE_CASE						       =    nn.Parameter(lowerCAmelCase_  ,						requires_grad=old_value.requires_grad       )
															__SCREAMING_SNAKE_CASE						       =    new_value




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=False       ):




					'''simple docstring'''


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

										if (isinstance(lowerCAmelCase_  ,						nn.Linear       ) or isinstance(lowerCAmelCase_  ,						lowerCAmelCase_       )) 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(lowerCAmelCase_       ) for key in modules_to_not_convert       ):
																				with init_empty_weights():
																									if isinstance(lowerCAmelCase_  ,						lowerCAmelCase_       ):
																														__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    module.weight.shape
																									else:
																														__SCREAMING_SNAKE_CASE						       =    module.in_features
																														__SCREAMING_SNAKE_CASE						       =    module.out_features

																									if quantization_config.quantization_method() == "llm_int8":
																														__SCREAMING_SNAKE_CASE						       =    bnb.nn.LinearabitLt(
																														    lowerCAmelCase_  ,						lowerCAmelCase_  ,						module.bias is not None  ,						has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight  ,						threshold=quantization_config.llm_inta_threshold  ,						)
																														__SCREAMING_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:
																																			__SCREAMING_SNAKE_CASE						       =    bnb.nn.Linearabit(
																																			    lowerCAmelCase_  ,						lowerCAmelCase_  ,						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  ,						)
																																			__SCREAMING_SNAKE_CASE						       =    True
                    # Store the module class in case we need to transpose the weight later
																									__SCREAMING_SNAKE_CASE						       =    type(lowerCAmelCase_       )
																									# Force requires grad to False to avoid unexpected errors
																									model._modules[name].requires_grad_(lowerCAmelCase_       )
										if len(list(module.children()       )       ) > 0:
															__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    _replace_with_bnb_linear(
															    lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						has_been_replaced=lowerCAmelCase_  ,						)
										# Remove the last key for recursion
										current_key_name.pop(-1       )
					return model, has_been_replaced




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None  ,						lowerCAmelCase_=None       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
					__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    _replace_with_bnb_linear(
					    lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_       )

					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        UpperCAmelCase__	(*lowerCAmelCase_  ,						**lowerCAmelCase_       ):




					'''simple docstring'''


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




def        UpperCAmelCase__	(*lowerCAmelCase_  ,						**lowerCAmelCase_       ):




					'''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"  ,						lowerCAmelCase_  ,						)
					return set_module_quantized_tensor_to_device(*lowerCAmelCase_  ,						**lowerCAmelCase_       )




def        UpperCAmelCase__	(lowerCAmelCase_       ):




					'''simple docstring'''


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

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

					# Check if it is a base model
					__SCREAMING_SNAKE_CASE						       =    not hasattr(lowerCAmelCase_  ,						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
					__SCREAMING_SNAKE_CASE						       =    list(model.named_children()       )
					__SCREAMING_SNAKE_CASE						       =    [list_modules[-1][0]]

					# add last module together with tied weights
					__SCREAMING_SNAKE_CASE						       =    set(lowerCAmelCase_       ) - set(lowerCAmelCase_       )
					__SCREAMING_SNAKE_CASE						       =    list(set(lowerCAmelCase_       )       ) + list(lowerCAmelCase_       )

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

					return filtered_module_names




"""simple docstring"""







import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


a__    :							Union[str, Any]            =					logging.get_logger(__name__)




def        UpperCAmelCase__	(lowerCAmelCase_       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    R"\w+[.]\d+"
					__SCREAMING_SNAKE_CASE						       =    re.findall(lowerCAmelCase_  ,						lowerCAmelCase_       )
					for pat in pats:
										__SCREAMING_SNAKE_CASE						       =    key.replace(lowerCAmelCase_  ,						"_".join(pat.split("."       )       )       )
					return key




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("scale",)
					if (
					    any("norm" in str_ for str_ in pt_tuple_key       )
					    and (pt_tuple_key[-1] == "bias")
					    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
					    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
					):
										__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("scale",)
										return renamed_pt_tuple_key, pt_tensor
					elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
										__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("scale",)
										return renamed_pt_tuple_key, pt_tensor

					# embedding
					if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
										__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("embedding",)
										return renamed_pt_tuple_key, pt_tensor

					# conv layer
					__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("kernel",)
					if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
										__SCREAMING_SNAKE_CASE						       =    pt_tensor.transpose(2  ,						3  ,						1  ,						0       )
										return renamed_pt_tuple_key, pt_tensor

					# linear layer
					__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("kernel",)
					if pt_tuple_key[-1] == "weight":
										__SCREAMING_SNAKE_CASE						       =    pt_tensor.T
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm weight
					__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("weight",)
					if pt_tuple_key[-1] == "gamma":
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm bias
					__SCREAMING_SNAKE_CASE						       =    pt_tuple_key[:-1] + ("bias",)
					if pt_tuple_key[-1] == "beta":
										return renamed_pt_tuple_key, pt_tensor

					return pt_tuple_key, pt_tensor




def        UpperCAmelCase__	(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_=42       ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE						       =    {k: v.numpy() for k, v in pt_state_dict.items()}

					# Step 2: Since the model is stateless, get random Flax params
					__SCREAMING_SNAKE_CASE						       =    flax_model.init_weights(PRNGKey(lowerCAmelCase_       )       )

					__SCREAMING_SNAKE_CASE						       =    flatten_dict(lowerCAmelCase_       )
					__SCREAMING_SNAKE_CASE						       =    {}

					# Need to change some parameters name to match Flax names
					for pt_key, pt_tensor in pt_state_dict.items():
										__SCREAMING_SNAKE_CASE						       =    rename_key(lowerCAmelCase_       )
										__SCREAMING_SNAKE_CASE						       =    tuple(renamed_pt_key.split("."       )       )

										# Correctly rename weight parameters
										__SCREAMING_SNAKE_CASE       ,  __SCREAMING_SNAKE_CASE						       =    rename_key_and_reshape_tensor(lowerCAmelCase_  ,						lowerCAmelCase_  ,						lowerCAmelCase_       )

										if flax_key in random_flax_state_dict:
															if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																				raise ValueError(
																				    f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																				    f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""       )

        # also add unexpected weight so that warning is thrown
										__SCREAMING_SNAKE_CASE						       =    jnp.asarray(lowerCAmelCase_       )

					return unflatten_dict(lowerCAmelCase_       )
