






"""simple docstring"""


import torch
from torch import nn




class 						__UpperCamelCase     (     nn.Module						):
					def __init__(							self							:	Optional[int]     , UpperCAmelCase							:	Any     , UpperCAmelCase							:	Optional[int]     , UpperCAmelCase							:	List[Any]     , UpperCAmelCase							:	str     , UpperCAmelCase							:	List[Any]=1     , UpperCAmelCase							:	Any=False   )	->       str:
											super().__init__()

											lowerCAmelCase   :int						=			n_token
											lowerCAmelCase   :str						=			d_embed
											lowerCAmelCase   :str						=			d_proj

											lowerCAmelCase   :Optional[int]						=			cutoffs + [n_token]
											lowerCAmelCase   :Dict						=			[0] + self.cutoffs
											lowerCAmelCase   :Dict						=			div_val

											lowerCAmelCase   :Dict						=			self.cutoffs[0]
											lowerCAmelCase   :Tuple						=			len(self.cutoffs   ) - 1
											lowerCAmelCase   :int						=			self.shortlist_size + self.n_clusters

											if self.n_clusters > 0:
																	lowerCAmelCase   :List[str]						=			nn.Parameter(torch.zeros(self.n_clusters     , self.d_embed   )   )
																	lowerCAmelCase   :Union[str, Any]						=			nn.Parameter(torch.zeros(self.n_clusters   )   )

											lowerCAmelCase   :str						=			nn.ModuleList()
											lowerCAmelCase   :Dict						=			nn.ParameterList()

											if div_val == 1:
																	for i in range(len(self.cutoffs   )   ):
																							if d_proj != d_embed:
																													self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCAmelCase     , UpperCAmelCase   )   )   )
																							else:
																													self.out_projs.append(UpperCAmelCase   )

																	self.out_layers.append(nn.Linear(UpperCAmelCase     , UpperCAmelCase   )   )
											else:
																	for i in range(len(self.cutoffs   )   ):
																							lowerCAmelCase	, lowerCAmelCase   :int						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowerCAmelCase   :Optional[int]						=			d_embed // (div_val**i)

																							self.out_projs.append(nn.Parameter(torch.FloatTensor(UpperCAmelCase     , UpperCAmelCase   )   )   )

																							self.out_layers.append(nn.Linear(UpperCAmelCase     , r_idx - l_idx   )   )

											lowerCAmelCase   :Union[str, Any]						=			keep_order
					def        UpperCAmelCase__ (							self							:	str     , UpperCAmelCase							:	int     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	int     , UpperCAmelCase							:	Tuple   )	->       Dict:
											if proj is None:
																	lowerCAmelCase   :List[Any]						=			nn.functional.linear(UpperCAmelCase     , UpperCAmelCase     , bias=UpperCAmelCase   )
											else:
																	# if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
																	lowerCAmelCase   :Union[str, Any]						=			nn.functional.linear(UpperCAmelCase     , proj.t().contiguous()   )
																	lowerCAmelCase   :Optional[int]						=			nn.functional.linear(UpperCAmelCase     , UpperCAmelCase     , bias=UpperCAmelCase   )
																	# else:
																	#     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
																	#     if bias is not None:
																	#         logit = logit + bias

											return logit
					def        UpperCAmelCase__ (							self							:	Dict     , UpperCAmelCase							:	Optional[Any]     , UpperCAmelCase							:	Union[str, Any]=None     , UpperCAmelCase							:	Optional[int]=False   )	->       Tuple:

											if labels is not None:
																	# Shift so that tokens < n predict n
																	lowerCAmelCase   :Optional[int]						=			hidden[..., :-1, :].contiguous()
																	lowerCAmelCase   :Optional[Any]						=			labels[..., 1:].contiguous()
																	lowerCAmelCase   :Any						=			hidden.view(-1     , hidden.size(-1   )   )
																	lowerCAmelCase   :int						=			labels.view(-1   )
																	if hidden.size(0   ) != labels.size(0   ):
																							raise RuntimeError('Input and labels should have the same size in the batch dimension.'   )
											else:
																	lowerCAmelCase   :List[Any]						=			hidden.view(-1     , hidden.size(-1   )   )

											if self.n_clusters == 0:
																	lowerCAmelCase   :Optional[Any]						=			self._compute_logit(UpperCAmelCase     , self.out_layers[0].weight     , self.out_layers[0].bias     , self.out_projs[0]   )
																	if labels is not None:
																							lowerCAmelCase   :Dict						=			labels != -100
																							lowerCAmelCase   :List[Any]						=			torch.zeros_like(UpperCAmelCase     , dtype=hidden.dtype     , device=hidden.device   )
																							lowerCAmelCase   :int						=			(
																							    -nn.functional.log_softmax(UpperCAmelCase     , dim=-1   )[mask].gather(1     , labels[mask].unsqueeze(1   )   ).squeeze(1   )
																							)
																	else:
																							lowerCAmelCase   :Dict						=			nn.functional.log_softmax(UpperCAmelCase     , dim=-1   )
											else:
																	# construct weights and biases
																	lowerCAmelCase	, lowerCAmelCase   :str						=			[], []
																	for i in range(len(self.cutoffs   )   ):
																							if self.div_val == 1:
																													lowerCAmelCase	, lowerCAmelCase   :Tuple						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
																													lowerCAmelCase   :Union[str, Any]						=			self.out_layers[0].weight[l_idx:r_idx]
																													lowerCAmelCase   :Optional[Any]						=			self.out_layers[0].bias[l_idx:r_idx]
																							else:
																													lowerCAmelCase   :Dict						=			self.out_layers[i].weight
																													lowerCAmelCase   :List[Any]						=			self.out_layers[i].bias

																							if i == 0:
																													lowerCAmelCase   :Dict						=			torch.cat([weight_i, self.cluster_weight]     , dim=0   )
																													lowerCAmelCase   :List[str]						=			torch.cat([bias_i, self.cluster_bias]     , dim=0   )

																							weights.append(UpperCAmelCase   )
																							biases.append(UpperCAmelCase   )

																	lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Optional[Any]						=			weights[0], biases[0], self.out_projs[0]

																	lowerCAmelCase   :List[Any]						=			self._compute_logit(UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase   )
																	lowerCAmelCase   :List[Any]						=			nn.functional.log_softmax(UpperCAmelCase     , dim=1   )

																	if labels is None:
																							lowerCAmelCase   :Optional[int]						=			hidden.new_empty((head_logit.size(0   ), self.n_token)   )
																	else:
																							lowerCAmelCase   :List[Any]						=			torch.zeros_like(UpperCAmelCase     , dtype=hidden.dtype     , device=hidden.device   )

																	lowerCAmelCase   :List[Any]						=			0
																	lowerCAmelCase   :List[str]						=			[0] + self.cutoffs
																	for i in range(len(UpperCAmelCase   ) - 1   ):
																							lowerCAmelCase	, lowerCAmelCase   :List[Any]						=			cutoff_values[i], cutoff_values[i + 1]

																							if labels is not None:
																													lowerCAmelCase   :Any						=			(labels >= l_idx) & (labels < r_idx)
																													lowerCAmelCase   :Union[str, Any]						=			mask_i.nonzero().squeeze()

																													if indices_i.numel() == 0:
																																			continue

																													lowerCAmelCase   :Optional[int]						=			labels.index_select(0     , UpperCAmelCase   ) - l_idx
																													lowerCAmelCase   :Dict						=			head_logprob.index_select(0     , UpperCAmelCase   )
																													lowerCAmelCase   :List[str]						=			hidden.index_select(0     , UpperCAmelCase   )
																							else:
																													lowerCAmelCase   :Optional[int]						=			hidden

																							if i == 0:
																													if labels is not None:
																																			lowerCAmelCase   :Union[str, Any]						=			head_logprob_i.gather(1     , target_i[:, None]   ).squeeze(1   )
																													else:
																																			lowerCAmelCase   :int						=			head_logprob[:, : self.cutoffs[0]]
																							else:
																													lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :List[str]						=			weights[i], biases[i], self.out_projs[i]

																													lowerCAmelCase   :Optional[Any]						=			self._compute_logit(UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase   )
																													lowerCAmelCase   :Optional[int]						=			nn.functional.log_softmax(UpperCAmelCase     , dim=1   )
																													lowerCAmelCase   :str						=			self.cutoffs[0] + i - 1  # No probability for the head cluster
																													if labels is not None:
																																			lowerCAmelCase   :Optional[int]						=			head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
																																			    1     , target_i[:, None]   ).squeeze(1   )
																													else:
																																			lowerCAmelCase   :Any						=			head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
																																			lowerCAmelCase   :Any						=			logprob_i

																							if labels is not None:
																													if (hasattr(self     , 'keep_order'   ) and self.keep_order) or keep_order:
																																			out.index_copy_(0     , UpperCAmelCase     , -logprob_i   )
																													else:
																																			out[offset : offset + logprob_i.size(0   )].copy_(-logprob_i   )
																													offset += logprob_i.size(0   )

											return out
					def        UpperCAmelCase__ (							self							:	List[Any]     , UpperCAmelCase							:	Optional[Any]   )	->       List[str]:
											if self.n_clusters == 0:
																	lowerCAmelCase   :List[Any]						=			self._compute_logit(UpperCAmelCase     , self.out_layers[0].weight     , self.out_layers[0].bias     , self.out_projs[0]   )
																	return nn.functional.log_softmax(UpperCAmelCase     , dim=-1   )
											else:
																	# construct weights and biases
																	lowerCAmelCase	, lowerCAmelCase   :str						=			[], []
																	for i in range(len(self.cutoffs   )   ):
																							if self.div_val == 1:
																													lowerCAmelCase	, lowerCAmelCase   :Optional[int]						=			self.cutoff_ends[i], self.cutoff_ends[i + 1]
																													lowerCAmelCase   :Dict						=			self.out_layers[0].weight[l_idx:r_idx]
																													lowerCAmelCase   :List[Any]						=			self.out_layers[0].bias[l_idx:r_idx]
																							else:
																													lowerCAmelCase   :List[str]						=			self.out_layers[i].weight
																													lowerCAmelCase   :int						=			self.out_layers[i].bias

																							if i == 0:
																													lowerCAmelCase   :Dict						=			torch.cat([weight_i, self.cluster_weight]     , dim=0   )
																													lowerCAmelCase   :Optional[int]						=			torch.cat([bias_i, self.cluster_bias]     , dim=0   )

																							weights.append(UpperCAmelCase   )
																							biases.append(UpperCAmelCase   )

																	lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Union[str, Any]						=			weights[0], biases[0], self.out_projs[0]
																	lowerCAmelCase   :Union[str, Any]						=			self._compute_logit(UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase   )

																	lowerCAmelCase   :Tuple						=			hidden.new_empty((head_logit.size(0   ), self.n_token)   )
																	lowerCAmelCase   :Union[str, Any]						=			nn.functional.log_softmax(UpperCAmelCase     , dim=1   )

																	lowerCAmelCase   :List[Any]						=			[0] + self.cutoffs
																	for i in range(len(UpperCAmelCase   ) - 1   ):
																							lowerCAmelCase	, lowerCAmelCase   :Union[str, Any]						=			cutoff_values[i], cutoff_values[i + 1]

																							if i == 0:
																													lowerCAmelCase   :Dict						=			head_logprob[:, : self.cutoffs[0]]
																							else:
																													lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Any						=			weights[i], biases[i], self.out_projs[i]

																													lowerCAmelCase   :Optional[Any]						=			self._compute_logit(UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase   )
																													lowerCAmelCase   :Dict						=			nn.functional.log_softmax(UpperCAmelCase     , dim=1   )

																													lowerCAmelCase   :Dict						=			head_logprob[:, -i] + tail_logprob_i
																													lowerCAmelCase   :Any						=			logprob_i

																	return out







"""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()