



"""simple docstring"""

import tensorflow as tf

from ...tf_utils import shape_list




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

 def __init__(     self		:       int  ,     lowercase_		:       List[str]  ,     lowercase_		:       Optional[Any]  ,     lowercase_		:       Optional[int]  ,     lowercase_		:       Optional[Any]  ,     lowercase_		:       Union[str, Any]=1  ,     lowercase_		:       List[str]=False  ,     **lowercase_		:       int				):
     super().__init__(**_A				)

     snake_case_ :     Union[str, Any]				=	vocab_size
     snake_case_ :     Dict				=	d_embed
     snake_case_ :     Optional[Any]				=	d_proj

     snake_case_ :     Union[str, Any]				=	cutoffs + [vocab_size]
     snake_case_ :     List[str]				=	[0] + self.cutoffs
     snake_case_ :     List[str]				=	div_val

     snake_case_ :     int				=	self.cutoffs[0]
     snake_case_ :     Any				=	len(self.cutoffs				) - 1
     snake_case_ :     int				=	self.shortlist_size + self.n_clusters
     snake_case_ :     Optional[int]				=	keep_order

     snake_case_ :     Tuple				=	[]
     snake_case_ :     Union[str, Any]				=	[]

 def 						_snake_case					(     self		:       Dict  ,     lowercase_		:       Any				):


     if self.n_clusters > 0:
         snake_case_ :     Tuple				=	self.add_weight(
             shape=(self.n_clusters, self.d_embed)  ,     initializer='''zeros'''  ,     trainable=_A  ,     name='''cluster_weight'''				)
         snake_case_ :     Optional[Any]				=	self.add_weight(
             shape=(self.n_clusters,)  ,     initializer='''zeros'''  ,     trainable=_A  ,     name='''cluster_bias'''				)

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

             snake_case_ :     List[Any]				=	self.add_weight(
                 shape=(d_emb_i, self.d_proj)  ,     initializer='''zeros'''  ,     trainable=_A  ,     name=f"out_projs_._{i}"				)
             self.out_projs.append(_A				)
             snake_case_ :     Optional[int]				=	self.add_weight(
                 shape=(r_idx - l_idx, d_emb_i)  ,     initializer='''zeros'''  ,     trainable=_A  ,     name=f"out_layers_._{i}_._weight"  ,     )
             snake_case_ :     List[Any]				=	self.add_weight(
                 shape=(r_idx - l_idx,)  ,     initializer='''zeros'''  ,     trainable=_A  ,     name=f"out_layers_._{i}_._bias"  ,     )
             self.out_layers.append((weight, bias)				)
     super().build(_A				)

 @staticmethod
 def 						_snake_case					(     lowercase_		:       Tuple  ,     lowercase_		:       Optional[Any]  ,     lowercase_		:       int  ,     lowercase_		:       Tuple=None				):
     snake_case_ :     str				=	x
     if proj is not None:
         snake_case_ :     Any				=	tf.einsum('''ibd,ed->ibe'''  ,     _A  ,     _A				)
     return tf.einsum('''ibd,nd->ibn'''  ,     _A  ,     _A				) + b

 @staticmethod
 def 						_snake_case					(     lowercase_		:       Tuple  ,     lowercase_		:       str				):
     snake_case_ :     List[Any]				=	shape_list(_A				)
     snake_case_ :     Any				=	tf.range(lp_size[0]  ,     dtype=target.dtype				)
     snake_case_ :     Dict				=	tf.stack([r, target]  ,     1				)
     return tf.gather_nd(_A  ,     _A				)






 def 						_snake_case					(     self		:       List[Any]  ,     lowercase_		:       Union[str, Any]  ,     lowercase_		:       Union[str, Any]  ,     lowercase_		:       List[Any]=True  ,     lowercase_		:       Any=False				):
     snake_case_ :     Tuple				=	0
     if self.n_clusters == 0:
         snake_case_ :     int				=	self._logit(_A  ,     self.out_layers[0][0]  ,     self.out_layers[0][1]  ,     self.out_projs[0]				)
         if target is not None:
             snake_case_ :     Dict				=	tf.nn.sparse_softmax_cross_entropy_with_logits(labels=_A  ,     logits=_A				)
         snake_case_ :     List[Any]				=	tf.nn.log_softmax(_A  ,     axis=-1				)
     else:
         snake_case_ :     Dict				=	shape_list(_A				)
         snake_case_ :     Optional[int]				=	[]
         snake_case_ :     Optional[int]				=	tf.zeros(hidden_sizes[:2]				)
         for i in range(len(self.cutoffs				)				):
             snake_case_ :     Optional[Any]				=	self.cutoff_ends[i], self.cutoff_ends[i + 1]
             if target is not None:
                 snake_case_ :     List[str]				=	(target >= l_idx) & (target < r_idx)
                 snake_case_ :     Optional[int]				=	tf.where(_A				)
                 snake_case_ :     Union[str, Any]				=	tf.boolean_mask(_A  ,     _A				) - l_idx

             if self.div_val == 1:
                 snake_case_ :     Any				=	self.out_layers[0][0][l_idx:r_idx]
                 snake_case_ :     Any				=	self.out_layers[0][1][l_idx:r_idx]
             else:
                 snake_case_ :     Tuple				=	self.out_layers[i][0]
                 snake_case_ :     List[str]				=	self.out_layers[i][1]

             if i == 0:
                 snake_case_ :     List[str]				=	tf.concat([cur_W, self.cluster_weight]  ,     0				)
                 snake_case_ :     Any				=	tf.concat([cur_b, self.cluster_bias]  ,     0				)

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

     if target is not None:
         if return_mean:
             snake_case_ :     List[str]				=	tf.reduce_mean(_A				)
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(_A				)

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

     return out









'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase__    					=      {
    '''configuration_distilbert''': [
        '''DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''DistilBertConfig''',
        '''DistilBertOnnxConfig''',
    ],
    '''tokenization_distilbert''': ['''DistilBertTokenizer'''],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      ['''DistilBertTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''DistilBertForMaskedLM''',
      '''DistilBertForMultipleChoice''',
      '''DistilBertForQuestionAnswering''',
      '''DistilBertForSequenceClassification''',
      '''DistilBertForTokenClassification''',
      '''DistilBertModel''',
      '''DistilBertPreTrainedModel''',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFDistilBertForMaskedLM''',
      '''TFDistilBertForMultipleChoice''',
      '''TFDistilBertForQuestionAnswering''',
      '''TFDistilBertForSequenceClassification''',
      '''TFDistilBertForTokenClassification''',
      '''TFDistilBertMainLayer''',
      '''TFDistilBertModel''',
      '''TFDistilBertPreTrainedModel''',
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__    					=      [
      '''FlaxDistilBertForMaskedLM''',
      '''FlaxDistilBertForMultipleChoice''',
      '''FlaxDistilBertForQuestionAnswering''',
      '''FlaxDistilBertForSequenceClassification''',
      '''FlaxDistilBertForTokenClassification''',
      '''FlaxDistilBertModel''',
      '''FlaxDistilBertPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_distilbert import (
      DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      DistilBertConfig,
      DistilBertOnnxConfig,
  )
  from .tokenization_distilbert import DistilBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_distilbert_fast import DistilBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_distilbert import (
        DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DistilBertForMaskedLM,
        DistilBertForMultipleChoice,
        DistilBertForQuestionAnswering,
        DistilBertForSequenceClassification,
        DistilBertForTokenClassification,
        DistilBertModel,
        DistilBertPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_distilbert import (
        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDistilBertForMaskedLM,
        TFDistilBertForMultipleChoice,
        TFDistilBertForQuestionAnswering,
        TFDistilBertForSequenceClassification,
        TFDistilBertForTokenClassification,
        TFDistilBertMainLayer,
        TFDistilBertModel,
        TFDistilBertPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_distilbert import (
        FlaxDistilBertForMaskedLM,
        FlaxDistilBertForMultipleChoice,
        FlaxDistilBertForQuestionAnswering,
        FlaxDistilBertForSequenceClassification,
        FlaxDistilBertForTokenClassification,
        FlaxDistilBertModel,
        FlaxDistilBertPreTrainedModel,
    )

else:
  import sys

  UpperCamelCase__    					=      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
