






'''simple docstring'''







import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


UpperCamelCase     :   int	      =		logging.get_logger(__name__)
def     A__      (					__lowerCAmelCase		:      int ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      List[Any] ,     __lowerCAmelCase		:      Optional[Any]=False    ):
						try:
												import torch  # noqa: F401
						except ImportError:
												logger.error(
												    """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
												    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
												    """ instructions."""    )
												raise

						if not is_sharded:
												lowerCamelCase__								=	os.path.abspath(__lowerCAmelCase    )
												logger.info(F'''Loading PyTorch weights from {pt_path}'''    )

												lowerCamelCase__								=	torch.load(__lowerCAmelCase ,     map_location="""cpu"""    )
												logger.info(F'''PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()    ):,} parameters.'''    )

												lowerCamelCase__								=	convert_pytorch_state_dict_to_flax(__lowerCAmelCase ,     __lowerCAmelCase    )
						else:
												# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
												lowerCamelCase__								=	convert_pytorch_sharded_state_dict_to_flax(__lowerCAmelCase ,     __lowerCAmelCase    )
						return flax_state_dict
def     A__      (					__lowerCAmelCase		:      Tuple[str] ,     __lowerCAmelCase		:      np.ndarray ,     __lowerCAmelCase		:      Dict[str, jnp.ndarray] ,     __lowerCAmelCase		:      str ,     ):

						def is_key_or_prefix_key_in_dict(__lowerCAmelCase		:      Tuple[str]    ) -> bool:
												return len(set(__lowerCAmelCase    ) & {key, (model_prefix,) + key}    ) > 0

						# layer norm
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""scale""",)
						if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# batch norm layer mean
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""mean""",)
						if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# batch norm layer var
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""var""",)
						if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

						# embedding
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""embedding""",)
						if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												return renamed_pt_tuple_key, pt_tensor

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

						# linear layer
						lowerCamelCase__								=	pt_tuple_key[:-1] + ("""kernel""",)
						if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase    ):
												lowerCamelCase__								=	pt_tensor.T
												return renamed_pt_tuple_key, pt_tensor

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

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

						# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
						lowerCamelCase__								=	None
						if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
												lowerCamelCase__								=	pt_tuple_key[-2] + """_g"""
						elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
												lowerCamelCase__								=	pt_tuple_key[-2] + """_v"""
						if name is not None:
												lowerCamelCase__								=	pt_tuple_key[:-3] + (name,)
												return renamed_pt_tuple_key, pt_tensor

						return pt_tuple_key, pt_tensor
def     A__      (					__lowerCAmelCase		:      Tuple ,     __lowerCAmelCase		:      Any    ):
						# convert pytorch tensor to numpy
						lowerCamelCase__								=	{k: v.numpy() for k, v in pt_state_dict.items()}

						lowerCamelCase__								=	flax_model.base_model_prefix

						# use params dict if the model contains batch norm layers
						if "params" in flax_model.params:
												lowerCamelCase__								=	flax_model.params["""params"""]
						else:
												lowerCamelCase__								=	flax_model.params
						lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )

						# add batch_stats keys,values to dict
						if "batch_stats" in flax_model.params:
												lowerCamelCase__								=	flatten_dict(flax_model.params["""batch_stats"""]    )
												random_flax_state_dict.update(__lowerCAmelCase    )

						lowerCamelCase__								=	{}

						lowerCamelCase__								=	(model_prefix not in flax_model_params) and (
						    model_prefix in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
						)
						lowerCamelCase__								=	(model_prefix in flax_model_params) and (
						    model_prefix not in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
						)

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

												# remove base model prefix if necessary
												lowerCamelCase__								=	pt_tuple_key[0] == model_prefix
												if load_model_with_head_into_base_model and has_base_model_prefix:
																		lowerCamelCase__								=	pt_tuple_key[1:]

												# Correctly rename weight parameters
												lowerCamelCase__							,      lowerCamelCase__								=	rename_key_and_reshape_tensor(
												    __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )

												# add model prefix if necessary
												lowerCamelCase__								=	(model_prefix,) + flax_key in random_flax_state_dict
												if load_base_model_into_model_with_head and require_base_model_prefix:
																		lowerCamelCase__								=	(model_prefix,) + flax_key

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

        # add batch stats if the model contains batchnorm layers
												if "batch_stats" in flax_model.params:
																		if "mean" in flax_key[-1] or "var" in flax_key[-1]:
																								lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																								continue
																		# remove num_batches_tracked key
																		if "num_batches_tracked" in flax_key[-1]:
																								flax_state_dict.pop(__lowerCAmelCase ,     __lowerCAmelCase    )
																								continue

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

												else:
																		# also add unexpected weight so that warning is thrown
																		lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )

						return unflatten_dict(__lowerCAmelCase    )
def     A__      (					__lowerCAmelCase		:      Optional[int] ,     __lowerCAmelCase		:      List[str]    ):
						import torch

						# Load the index
						lowerCamelCase__								=	{}
						for shard_file in shard_filenames:
												# load using msgpack utils
												lowerCamelCase__								=	torch.load(__lowerCAmelCase    )
												lowerCamelCase__								=	{k: v.numpy() for k, v in pt_state_dict.items()}

												lowerCamelCase__								=	flax_model.base_model_prefix

												# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
												if "batch_stats" in flax_model.params:
																		lowerCamelCase__								=	flax_model.params["""params"""]

																		lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )
																		random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]    )    )
												else:
																		lowerCamelCase__								=	flax_model.params
																		lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )

												lowerCamelCase__								=	(model_prefix not in flax_model_params) and (
												    model_prefix in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
												)
												lowerCamelCase__								=	(model_prefix in flax_model_params) and (
												    model_prefix not in {k.split("""."""    )[0] for k in pt_state_dict.keys()}
												)
												# Need to change some parameters name to match Flax names
												for pt_key, pt_tensor in pt_state_dict.items():
																		lowerCamelCase__								=	tuple(pt_key.split("""."""    )    )

																		# remove base model prefix if necessary
																		lowerCamelCase__								=	pt_tuple_key[0] == model_prefix
																		if load_model_with_head_into_base_model and has_base_model_prefix:
																								lowerCamelCase__								=	pt_tuple_key[1:]

																		# Correctly rename weight parameters
																		lowerCamelCase__							,      lowerCamelCase__								=	rename_key_and_reshape_tensor(
																		    __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase ,     __lowerCAmelCase    )
																		# add model prefix if necessary
																		lowerCamelCase__								=	(model_prefix,) + flax_key in random_flax_state_dict
																		if load_base_model_into_model_with_head and require_base_model_prefix:
																								lowerCamelCase__								=	(model_prefix,) + flax_key

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

            # add batch stats if the model contains batchnorm layers
																		if "batch_stats" in flax_model.params:
																								if "mean" in flax_key[-1]:
																														lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																														continue
																								if "var" in flax_key[-1]:
																														lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
																														continue
																								# remove num_batches_tracked key
																								if "num_batches_tracked" in flax_key[-1]:
																														flax_state_dict.pop(__lowerCAmelCase ,     __lowerCAmelCase    )
																														continue

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

																		else:
																								# also add unexpected weight so that warning is thrown
																								lowerCamelCase__								=	jnp.asarray(__lowerCAmelCase    )
						return unflatten_dict(__lowerCAmelCase    )
def     A__      (					__lowerCAmelCase		:      Any ,     __lowerCAmelCase		:      Tuple    ):
						lowerCamelCase__								=	os.path.abspath(__lowerCAmelCase    )
						logger.info(F'''Loading Flax weights from {flax_checkpoint_path}'''    )

						# import correct flax class
						lowerCamelCase__								=	getattr(__lowerCAmelCase ,     """Flax""" + model.__class__.__name__    )

						# load flax weight dict
						with open(__lowerCAmelCase ,     """rb"""    ) as state_f:
												try:
																		lowerCamelCase__								=	from_bytes(__lowerCAmelCase ,     state_f.read()    )
												except UnpicklingError:
																		raise EnvironmentError(F'''Unable to convert {flax_checkpoint_path} to Flax deserializable object. '''    )

						return load_flax_weights_in_pytorch_model(__lowerCAmelCase ,     __lowerCAmelCase    )





def     A__      (					__lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Tuple    ):

						try:
												import torch  # noqa: F401
						except ImportError:
												logger.error(
												    """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
												    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
												    """ instructions."""    )
												raise

						# check if we have bf16 weights
						lowerCamelCase__								=	flatten_dict(jax.tree_util.tree_map(lambda __lowerCAmelCase    : x.dtype == jnp.bfloataa ,     __lowerCAmelCase    )    ).values()
						if any(__lowerCAmelCase    ):
												# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
												# and bf16 is not fully supported in PT yet.
												logger.warning(
												    """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
												    """before loading those in PyTorch model."""    )
												lowerCamelCase__								=	jax.tree_util.tree_map(
												    lambda __lowerCAmelCase    : params.astype(np.floataa    ) if params.dtype == jnp.bfloataa else params ,     __lowerCAmelCase    )

						lowerCamelCase__								=	flatten_dict(__lowerCAmelCase    )
						lowerCamelCase__								=	pt_model.state_dict()

						lowerCamelCase__								=	(pt_model.base_model_prefix in flax_state) and (
						    pt_model.base_model_prefix not in {k.split("""."""    )[0] for k in pt_model_dict.keys()}
						)
						lowerCamelCase__								=	(pt_model.base_model_prefix not in flax_state) and (
						    pt_model.base_model_prefix in {k.split("""."""    )[0] for k in pt_model_dict.keys()}
						)

						# keep track of unexpected & missing keys
						lowerCamelCase__								=	[]
						lowerCamelCase__								=	set(pt_model_dict.keys()    )

						for flax_key_tuple, flax_tensor in flax_state_dict.items():
												lowerCamelCase__								=	flax_key_tuple[0] == pt_model.base_model_prefix
												lowerCamelCase__								=	""".""".join((pt_model.base_model_prefix,) + flax_key_tuple    ) in pt_model_dict

												# adapt flax_key to prepare for loading from/to base model only
												if load_model_with_head_into_base_model and has_base_model_prefix:
																		lowerCamelCase__								=	flax_key_tuple[1:]
												elif load_base_model_into_model_with_head and require_base_model_prefix:
																		lowerCamelCase__								=	(pt_model.base_model_prefix,) + flax_key_tuple

												# rename flax weights to PyTorch format
												if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCAmelCase    ) not in pt_model_dict:
																		# conv layer
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)
																		lowerCamelCase__								=	jnp.transpose(__lowerCAmelCase ,     (3, 2, 0, 1)    )
												elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCAmelCase    ) not in pt_model_dict:
																		# linear layer
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)
																		lowerCamelCase__								=	flax_tensor.T
												elif flax_key_tuple[-1] in ["scale", "embedding"]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""weight""",)

												# adding batch stats from flax batch norm to pt
												elif "mean" in flax_key_tuple[-1]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""running_mean""",)
												elif "var" in flax_key_tuple[-1]:
																		lowerCamelCase__								=	flax_key_tuple[:-1] + ("""running_var""",)

												if "batch_stats" in flax_state:
																		lowerCamelCase__								=	""".""".join(flax_key_tuple[1:]    )  # Remove the params/batch_stats header
												else:
																		lowerCamelCase__								=	""".""".join(__lowerCAmelCase    )

												# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
												lowerCamelCase__								=	{}
												# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
												for key in pt_model_dict:
																		lowerCamelCase__								=	key.split("""."""    )
																		lowerCamelCase__								=	None
																		if key_components[-3::2] == ["parametrizations", "original0"]:
																								lowerCamelCase__								=	key_components[-2] + """_g"""
																		elif key_components[-3::2] == ["parametrizations", "original1"]:
																								lowerCamelCase__								=	key_components[-2] + """_v"""
																		if name is not None:
																								lowerCamelCase__								=	key_components[:-3] + [name]
																								lowerCamelCase__								=	""".""".join(__lowerCAmelCase    )
																								lowerCamelCase__								=	key

												if flax_key in special_pt_names:
																		lowerCamelCase__								=	special_pt_names[flax_key]

												if flax_key in pt_model_dict:
																		if flax_tensor.shape != pt_model_dict[flax_key].shape:
																								raise ValueError(
																								    F'''Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '''
																								    F'''to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'''    )
																		else:
																								# add weight to pytorch dict
																								lowerCamelCase__								=	np.asarray(__lowerCAmelCase    ) if not isinstance(__lowerCAmelCase ,     np.ndarray    ) else flax_tensor
																								lowerCamelCase__								=	torch.from_numpy(__lowerCAmelCase    )
																								# remove from missing keys
																								missing_keys.remove(__lowerCAmelCase    )
												else:
																		# weight is not expected by PyTorch model
																		unexpected_keys.append(__lowerCAmelCase    )

						pt_model.load_state_dict(__lowerCAmelCase    )

						# re-transform missing_keys to list
						lowerCamelCase__								=	list(__lowerCAmelCase    )

						if len(__lowerCAmelCase    ) > 0:
												logger.warning(
												    """Some weights of the Flax model were not used when initializing the PyTorch model"""
												    F''' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'''
												    F''' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'''
												    """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
												    F''' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'''
												    """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
												    """ FlaxBertForSequenceClassification model)."""    )
						else:
												logger.warning(F'''All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'''    )
						if len(__lowerCAmelCase    ) > 0:
												logger.warning(
												    F'''Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'''
												    F''' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'''
												    """ use it for predictions and inference."""    )
						else:
												logger.warning(
												    F'''All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'''
												    """If your task is similar to the task the model of the checkpoint was trained on, """
												    F'''you can already use {pt_model.__class__.__name__} for predictions without further training.'''    )

						return pt_model









"""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 						__UpperCamelCase     (     UpperCamelCase						):
					lowercase_    :				Optional[Any]      =  """Salesforce/blip-image-captioning-base"""
					lowercase_    :				Optional[int]      =  (
					    """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_    :				int      =  """image_captioner"""
					lowercase_    :				int      =  AutoModelForVisionaSeq

					lowercase_    :				List[str]      =  ["""image"""]
					lowercase_    :				int      =  ["""text"""]
					def __init__(							self							:	Any     , *UpperCAmelCase							:	str     , **UpperCAmelCase							:	List[str]   )	->       Union[str, Any]:
											requires_backends(self     , ['vision']   )
											super().__init__(*UpperCAmelCase     , **UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	int     , UpperCAmelCase							:	"Image"   )	->       Union[str, Any]:
											return self.pre_processor(images=UpperCAmelCase     , return_tensors='pt'   )
					def        UpperCAmelCase__ (							self							:	int     , UpperCAmelCase							:	Optional[int]   )	->       Optional[int]:
											return self.model.generate(**UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	Dict     , UpperCAmelCase							:	int   )	->       Dict:
											return self.pre_processor.batch_decode(UpperCAmelCase     , skip_special_tokens=UpperCAmelCase   )[0].strip()