<reponame>xiongma/bert2tf
from typing import Optional, Union, List, Tuple

import tensorflow as tf
from tensorflow.python.keras import backend as K

from .helper import get_initializer, shape_list


class BaseEmbedding(tf.keras.layers.Layer):
    """
    Base Embedding Layer
    """

    def __init__(self,
                 vocab_size: int,
                 embed_size: int,
                 hidden_dropout_prob: float = 0.1,
                 initializer_range: Optional[float] = None,
                 word_embedding_name: str = 'word_embeddings',
                 **kwargs):
        """

        :param vocab_size: vocab size
        :param embed_size: embedding size
        :param hidden_dropout_prob: hidden dropout probability
        :param initializer_range: word embedding weights initial range
        :param word_embedding_name: word embedding weights name
        """
        super(BaseEmbedding, self).__init__(**kwargs)

        self.vocab_size = vocab_size
        self.hidden_size = embed_size
        self.word_embeddings_name = word_embedding_name
        self.initializer_range = embed_size ** -0.5 if initializer_range is None else initializer_range

        self.dropout = tf.keras.layers.Dropout(hidden_dropout_prob)

    def build(self, input_shape):
        # Create and initialize weights. The random normal initializer was chosen
        # arbitrarily, and works well.
        self.word_embeddings = self.add_weight(
            self.word_embeddings_name,
            shape=[self.vocab_size, self.hidden_size],
            initializer=get_initializer(self.initializer_range)
        )

        super(BaseEmbedding, self).build(input_shape)

    def get_config(self):
        config = {
            'vocab_size': self.vocab_size,
            'hidden_size': self.hidden_size,
            'initializer_range': self.initializer_range,
        }
        base_config = super(BaseEmbedding, self).get_config()

        return dict(list(base_config.items()) + list(config.items()))

    def call(self, inputs: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]], mode: str = 'embedding',
             **kwargs) -> tf.Tensor:
        return getattr(self, mode)(inputs, **kwargs)

    def embedding(self,
                  inputs: Union[tf.Tensor, List[tf.Tensor], Tuple[tf.Tensor]],
                  training: tf.Tensor = False,
                  **kwargs) -> tf.Tensor:
        """
        Word embedding
        :param inputs: inputs
        :param training: training flag
        :return: word embeddings
        """
        embeds = self.get_embedding(inputs, **kwargs)
        embeds = self.dropout(embeds, training=training)

        return embeds

    def get_embedding(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
        """
        Word embedding main layer
        :param inputs: inputs
        :return: word embeddings
        """
        dtype = K.dtype(inputs)
        if dtype != 'int32' and dtype != 'int64':
            inputs = tf.cast(inputs, 'int32')

        return tf.gather(self.word_embeddings, inputs)

    def linear(self, inputs: tf.Tensor, **kwargs) -> tf.Tensor:
        """
        Map to hidden state to vocab
        :param inputs: inputs
        :return: sequence vocab distribution
        """
        batch_size = shape_list(inputs)[0]
        length = shape_list(inputs)[1]

        x = tf.reshape(inputs, [-1, self.hidden_size])
        logits = tf.matmul(x, self.word_embeddings, transpose_b=True)

        return tf.reshape(logits, [batch_size, length, self.vocab_size])
