



import tensorflow as tf

from ...tf_utils import shape_list



class 	lowerCamelCase      (						tf.keras.layers.Layer  ):







					'''simple docstring'''







					def __init__(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase=1    ,	lowerCAmelCase=False    ,	**lowerCAmelCase     ):
							super().__init__(**lowerCAmelCase     )

							UpperCAmelCase_						=				vocab_size
							UpperCAmelCase_						=				d_embed
							UpperCAmelCase_						=				d_proj

							UpperCAmelCase_						=				cutoffs + [vocab_size]
							UpperCAmelCase_						=				[0] + self.cutoffs
							UpperCAmelCase_						=				div_val

							UpperCAmelCase_						=				self.cutoffs[0]
							UpperCAmelCase_						=				len(self.cutoffs     ) - 1
							UpperCAmelCase_						=				self.shortlist_size + self.n_clusters
							UpperCAmelCase_						=				keep_order

							UpperCAmelCase_						=				[]
							UpperCAmelCase_						=				[]



					def    A__		(			self    ,	lowerCAmelCase     ):
							if self.n_clusters > 0:
									UpperCAmelCase_						=				self.add_weight(
									    shape=(self.n_clusters, self.d_embed)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name="cluster_weight"     )
									UpperCAmelCase_						=				self.add_weight(
									    shape=(self.n_clusters,)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name="cluster_bias"     )

							if self.div_val == 1:
									for i in range(len(self.cutoffs     )     ):
											if self.d_proj != self.d_embed:
													UpperCAmelCase_						=				self.add_weight(
													    shape=(self.d_embed, self.d_proj)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_projs_._{i}'''    ,	)
													self.out_projs.append(lowerCAmelCase     )
											else:
													self.out_projs.append(lowerCAmelCase     )
											UpperCAmelCase_						=				self.add_weight(
											    shape=(self.vocab_size, self.d_embed)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_layers_._{i}_._weight'''    ,	)
											UpperCAmelCase_						=				self.add_weight(
											    shape=(self.vocab_size,)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_layers_._{i}_._bias'''    ,	)
											self.out_layers.append((weight, bias)     )
							else:
									for i in range(len(self.cutoffs     )     ):
											UpperCAmelCase_					,						UpperCAmelCase_						=				self.cutoff_ends[i], self.cutoff_ends[i + 1]
											UpperCAmelCase_						=				self.d_embed // (self.div_val**i)

											UpperCAmelCase_						=				self.add_weight(
											    shape=(d_emb_i, self.d_proj)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_projs_._{i}'''     )
											self.out_projs.append(lowerCAmelCase     )
											UpperCAmelCase_						=				self.add_weight(
											    shape=(r_idx - l_idx, d_emb_i)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_layers_._{i}_._weight'''    ,	)
											UpperCAmelCase_						=				self.add_weight(
											    shape=(r_idx - l_idx,)    ,	initializer="zeros"    ,	trainable=lowerCAmelCase    ,	name=f'''out_layers_._{i}_._bias'''    ,	)
											self.out_layers.append((weight, bias)     )
							super().build(lowerCAmelCase     )



					@staticmethod
					def    A__		(			lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase=None     ):
							UpperCAmelCase_						=				x
							if proj is not None:
									UpperCAmelCase_						=				tf.einsum("ibd,ed->ibe"    ,	lowerCAmelCase    ,	lowerCAmelCase     )
							return tf.einsum("ibd,nd->ibn"    ,	lowerCAmelCase    ,	lowerCAmelCase     ) + b



					@staticmethod
					def    A__		(			lowerCAmelCase    ,	lowerCAmelCase     ):
							UpperCAmelCase_						=				shape_list(lowerCAmelCase     )
							UpperCAmelCase_						=				tf.range(lp_size[0]    ,	dtype=target.dtype     )
							UpperCAmelCase_						=				tf.stack([r, target]    ,	1     )
							return tf.gather_nd(lowerCAmelCase    ,	lowerCAmelCase     )



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase=True    ,	lowerCAmelCase=False     ):
							UpperCAmelCase_						=				0
							if self.n_clusters == 0:
									UpperCAmelCase_						=				self._logit(lowerCAmelCase    ,	self.out_layers[0][0]    ,	self.out_layers[0][1]    ,	self.out_projs[0]     )
									if target is not None:
											UpperCAmelCase_						=				tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCAmelCase    ,	logits=lowerCAmelCase     )
									UpperCAmelCase_						=				tf.nn.log_softmax(lowerCAmelCase    ,	axis=-1     )
							else:
									UpperCAmelCase_						=				shape_list(lowerCAmelCase     )
									UpperCAmelCase_						=				[]
									UpperCAmelCase_						=				tf.zeros(hidden_sizes[:2]     )
									for i in range(len(self.cutoffs     )     ):
											UpperCAmelCase_					,						UpperCAmelCase_						=				self.cutoff_ends[i], self.cutoff_ends[i + 1]
											if target is not None:
													UpperCAmelCase_						=				(target >= l_idx) & (target < r_idx)
													UpperCAmelCase_						=				tf.where(lowerCAmelCase     )
													UpperCAmelCase_						=				tf.boolean_mask(lowerCAmelCase    ,	lowerCAmelCase     ) - l_idx

											if self.div_val == 1:
													UpperCAmelCase_						=				self.out_layers[0][0][l_idx:r_idx]
													UpperCAmelCase_						=				self.out_layers[0][1][l_idx:r_idx]
											else:
													UpperCAmelCase_						=				self.out_layers[i][0]
													UpperCAmelCase_						=				self.out_layers[i][1]

											if i == 0:
													UpperCAmelCase_						=				tf.concat([cur_W, self.cluster_weight]    ,	0     )
													UpperCAmelCase_						=				tf.concat([cur_b, self.cluster_bias]    ,	0     )

													UpperCAmelCase_						=				self._logit(lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	self.out_projs[0]     )
													UpperCAmelCase_						=				tf.nn.log_softmax(lowerCAmelCase     )
													out.append(head_logprob[..., : self.cutoffs[0]]     )
													if target is not None:
															UpperCAmelCase_						=				tf.boolean_mask(lowerCAmelCase    ,	lowerCAmelCase     )
															UpperCAmelCase_						=				self._gather_logprob(lowerCAmelCase    ,	lowerCAmelCase     )
											else:
													UpperCAmelCase_						=				self._logit(lowerCAmelCase    ,	lowerCAmelCase    ,	lowerCAmelCase    ,	self.out_projs[i]     )
													UpperCAmelCase_						=				tf.nn.log_softmax(lowerCAmelCase     )
													UpperCAmelCase_						=				self.cutoffs[0] + i - 1  # No probability for the head cluster
													UpperCAmelCase_						=				head_logprob[..., cluster_prob_idx, None] + tail_logprob
													out.append(lowerCAmelCase     )
													if target is not None:
															UpperCAmelCase_						=				tf.boolean_mask(lowerCAmelCase    ,	lowerCAmelCase     )
															UpperCAmelCase_						=				tf.boolean_mask(lowerCAmelCase    ,	lowerCAmelCase     )
															UpperCAmelCase_						=				self._gather_logprob(lowerCAmelCase    ,	lowerCAmelCase     )
															cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
											if target is not None:
													loss += tf.scatter_nd(lowerCAmelCase    ,	-cur_logprob    ,	shape_list(lowerCAmelCase     )     )
									UpperCAmelCase_						=				tf.concat(lowerCAmelCase    ,	axis=-1     )

							if target is not None:
									if return_mean:
											UpperCAmelCase_						=				tf.reduce_mean(lowerCAmelCase     )
									# Add the training-time loss value to the layer using `self.add_loss()`.
									self.add_loss(lowerCAmelCase     )

									# Log the loss as a metric (we could log arbitrary metrics,
									# including different metrics for training and inference.
									self.add_metric(lowerCAmelCase    ,	name=self.name    ,	aggregation="mean" if return_mean else ""     )

							return out






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

SCREAMING_SNAKE_CASE     				=						logging.get_logger(__name__)



@add_end_docstrings(lowercase__  )
class 	lowerCamelCase      (						lowercase__  ):







					'''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    A__		(			self    ,	lowerCAmelCase=None    ,	lowerCAmelCase=None    ,	lowerCAmelCase=None     ):
							UpperCAmelCase_						=				{}
							UpperCAmelCase_						=				{}

							if prompt is not None:
									UpperCAmelCase_						=				prompt

							if generate_kwargs is not None:
									UpperCAmelCase_						=				generate_kwargs
							if max_new_tokens is not None:
									if "generate_kwargs" not in forward_kwargs:
											UpperCAmelCase_						=				{}
									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"     )
									UpperCAmelCase_						=				max_new_tokens
							return preprocess_params, forward_kwargs, {}



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



					def    A__		(			self    ,	lowerCAmelCase    ,	lowerCAmelCase=None     ):
							UpperCAmelCase_						=				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."     )

									UpperCAmelCase_						=				self.model.config.model_type

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

									elif model_type == "pix2struct":
											UpperCAmelCase_						=				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
											UpperCAmelCase_						=				self.image_processor(images=lowerCAmelCase    ,	return_tensors=self.framework     )
											UpperCAmelCase_						=				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:
									UpperCAmelCase_						=				self.image_processor(images=lowerCAmelCase    ,	return_tensors=self.framework     )

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

							return model_inputs



					def    A__		(			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"]     )
							):
									UpperCAmelCase_						=				None

							if generate_kwargs is None:
									UpperCAmelCase_						=				{}
							# 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.
							UpperCAmelCase_						=				model_inputs.pop(self.model.main_input_name     )
							UpperCAmelCase_						=				self.model.generate(lowerCAmelCase    ,	**lowerCAmelCase    ,	**lowerCAmelCase     )
							return model_outputs



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

