




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase__				     = {
    "configuration_lilt": ["LILT_PRETRAINED_CONFIG_ARCHIVE_MAP", "LiltConfig"],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			UpperCAmelCase__				     = [
			    "LILT_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "LiltForQuestionAnswering",
			    "LiltForSequenceClassification",
			    "LiltForTokenClassification",
			    "LiltModel",
			    "LiltPreTrainedModel",
			]

if TYPE_CHECKING:
			from .configuration_lilt import LILT_PRETRAINED_CONFIG_ARCHIVE_MAP, LiltConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_lilt import (
						    LILT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    LiltForQuestionAnswering,
						    LiltForSequenceClassification,
						    LiltForTokenClassification,
						    LiltModel,
						    LiltPreTrainedModel,
						)

else:
			import sys

			UpperCAmelCase__				     = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)












import tensorflow as tf

from ...tf_utils import shape_list







class a (     tf.keras.layers.Layer  ):
						"""simple docstring"""

						def __init__( self :     int				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Tuple=1				,							lowerCamelCase__ :     Optional[Any]=False				,							**lowerCamelCase__ :     Optional[int]				)    ->						str:
													"""simple docstring"""


													super().__init__(**lowerCamelCase__				)

													__lowercase      =    vocab_size
													__lowercase      =    d_embed
													__lowercase      =    d_proj

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

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

													__lowercase      =    []
													__lowercase      =    []






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				)    ->						Tuple:
													"""simple docstring"""


													if self.n_clusters > 0:
																				__lowercase      =    self.add_weight(
																				    shape=(self.n_clusters, self.d_embed)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name='''cluster_weight'''				)
																				__lowercase      =    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:
																																		__lowercase      =    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__				)
																											__lowercase      =    self.add_weight(
																											    shape=(self.vocab_size, self.d_embed)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    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				)				):
																											__lowercase     ,     __lowercase      =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
																											__lowercase      =    self.d_embed // (self.div_val**i)

																											__lowercase      =    self.add_weight(
																											    shape=(d_emb_i, self.d_proj)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_projs_._{i}'				)
																											self.out_projs.append(lowerCamelCase__				)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx, d_emb_i)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    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     UpperCAmelCase_      ( lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     int				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     List[str]=None				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    x
													if proj is not None:
																				__lowercase      =    tf.einsum('''ibd,ed->ibe'''				,							lowerCamelCase__				,							lowerCamelCase__				)
													return tf.einsum('''ibd,nd->ibn'''				,							lowerCamelCase__				,							lowerCamelCase__				) + b






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[str]				)    ->						str:
													"""simple docstring"""


													__lowercase      =    shape_list(lowerCamelCase__				)
													__lowercase      =    tf.range(lp_size[0]				,							dtype=target.dtype				)
													__lowercase      =    tf.stack([r, target]				,							1				)
													return tf.gather_nd(lowerCamelCase__				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Optional[int]=True				,							lowerCamelCase__ :     Optional[int]=False				)    ->						List[str]:
													"""simple docstring"""


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

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

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

																																		__lowercase      =    self._logit(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							self.out_projs[0]				)
																																		__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				)
																																		out.append(head_logprob[..., : self.cutoffs[0]]				)
																																		if target is not None:
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    self._gather_logprob(lowerCamelCase__				,							lowerCamelCase__				)
																											else:
																																		__lowercase      =    self._logit(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							self.out_projs[i]				)
																																		__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				)
																																		__lowercase      =    self.cutoffs[0] + i - 1  # No probability for the head cluster
																																		__lowercase      =    head_logprob[..., cluster_prob_idx, None] + tail_logprob
																																		out.append(lowerCamelCase__				)
																																		if target is not None:
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    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__				)				)
																				__lowercase      =    tf.concat(lowerCamelCase__				,							axis=-1				)

													if target is not None:
																				if return_mean:
																											__lowercase      =    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






