easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/layers/keras/transformer.py

@@ -0,0 +1,192 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import numpy as np
+import tensorflow as tf
+from tensorflow.python.keras.layers import Dense
+from tensorflow.python.keras.layers import Dropout
+from tensorflow.python.keras.layers import Embedding
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.layers.keras import MultiHeadAttention
+from easy_rec.python.layers.keras.layer_norm import LayerNormalization
+from easy_rec.python.layers.utils import Parameter
+from easy_rec.python.protos import seq_encoder_pb2
+
+
+class TransformerBlock(Layer):
+  """A transformer block combines multi-head attention and feed-forward networks with layer normalization and dropout.
+
+  Purpose: Combines attention and feed-forward layers with residual connections and normalization.
+  Components: Multi-head attention, feed-forward network, dropout, and layer normalization.
+  Output: Enhanced representation after applying attention and feed-forward layers.
+  """
+
+  def __init__(self, params, name='transformer_block', reuse=None, **kwargs):
+    super(TransformerBlock, self).__init__(name=name, **kwargs)
+    d_model = params.hidden_size
+    num_heads = params.num_attention_heads
+    mha_cfg = seq_encoder_pb2.MultiHeadAttention()
+    mha_cfg.num_heads = num_heads
+    mha_cfg.key_dim = d_model // num_heads
+    mha_cfg.dropout = params.get_or_default('attention_probs_dropout_prob', 0.0)
+    mha_cfg.return_attention_scores = False
+    args = Parameter.make_from_pb(mha_cfg)
+    self.mha = MultiHeadAttention(args, 'multi_head_attn')
+    dropout_rate = params.get_or_default('hidden_dropout_prob', 0.1)
+    ffn_units = params.get_or_default('intermediate_size', d_model)
+    ffn_act = params.get_or_default('hidden_act', 'relu')
+    self.ffn_dense1 = Dense(ffn_units, activation=ffn_act)
+    self.ffn_dense2 = Dense(d_model)
+    if tf.__version__ >= '2.0':
+      self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+      self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+    else:
+      self.layer_norm1 = LayerNormalization(epsilon=1e-6)
+      self.layer_norm2 = LayerNormalization(epsilon=1e-6)
+    self.dropout1 = Dropout(dropout_rate)
+    self.dropout2 = Dropout(dropout_rate)
+
+  def call(self, inputs, training=None, **kwargs):
+    x, mask = inputs
+    attn_output = self.mha([x, x, x], mask=mask, training=training)
+    attn_output = self.dropout1(attn_output, training=training)
+    out1 = self.layer_norm1(x + attn_output)
+    ffn_mid = self.ffn_dense1(out1)
+    ffn_output = self.ffn_dense2(ffn_mid)
+    ffn_output = self.dropout2(ffn_output, training=training)
+    out2 = self.layer_norm2(out1 + ffn_output)
+    return out2
+
+
+# Positional Encoding, https://www.tensorflow.org/text/tutorials/transformer
+def positional_encoding(length, depth):
+  depth = depth / 2
+  positions = np.arange(length)[:, np.newaxis]  # (seq, 1)
+  depths = np.arange(depth)[np.newaxis, :] / depth  # (1, depth)
+  angle_rates = 1 / (10000**depths)  # (1, depth)
+  angle_rads = positions * angle_rates  # (pos, depth)
+  pos_encoding = np.concatenate(
+      [np.sin(angle_rads), np.cos(angle_rads)], axis=-1)
+  return tf.cast(pos_encoding, dtype=tf.float32)
+
+
+class PositionalEmbedding(Layer):
+
+  def __init__(self, vocab_size, d_model, max_position, name='pos_embedding'):
+    super(PositionalEmbedding, self).__init__(name=name)
+    self.d_model = d_model
+    self.embedding = Embedding(vocab_size, d_model)
+    self.pos_encoding = positional_encoding(length=max_position, depth=d_model)
+
+  def call(self, x, training=None):
+    length = tf.shape(x)[1]
+    x = self.embedding(x)
+    # This factor sets the relative scale of the embedding and positional_encoding.
+    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
+    x = x + self.pos_encoding[tf.newaxis, :length, :]
+    return x
+
+
+class TransformerEncoder(Layer):
+  """The encoder consists of a stack of encoder layers.
+
+  It converts the input sequence into a set of embeddings enriched with positional information.
+  Purpose: Encodes the input sequence into a set of embeddings.
+  Components: Embedding layer, positional encoding, and a stack of transformer blocks.
+  Output: Encoded representation of the input sequence.
+  """
+
+  def __init__(self, params, name='transformer_encoder', reuse=None, **kwargs):
+    super(TransformerEncoder, self).__init__(name=name, **kwargs)
+    d_model = params.hidden_size
+    dropout_rate = params.get_or_default('hidden_dropout_prob', 0.1)
+    max_position = params.get_or_default('max_position_embeddings', 512)
+    num_layers = params.get_or_default('num_hidden_layers', 1)
+    vocab_size = params.vocab_size
+    logging.info('vocab size of TransformerEncoder(%s) is %d', name, vocab_size)
+    self.output_all = params.get_or_default('output_all_token_embeddings', True)
+    self.pos_encoding = PositionalEmbedding(vocab_size, d_model, max_position)
+    self.dropout = Dropout(dropout_rate)
+    self.enc_layers = [
+        TransformerBlock(params, 'layer_%d' % i) for i in range(num_layers)
+    ]
+    self._vocab_size = vocab_size
+    self._max_position = max_position
+
+  @property
+  def vocab_size(self):
+    return self._vocab_size
+
+  @property
+  def max_position(self):
+    return self._max_position
+
+  def call(self, inputs, training=None, **kwargs):
+    x, mask = inputs
+    # `x` is token-IDs shape: (batch, seq_len)
+    x = self.pos_encoding(x)  # Shape `(batch_size, seq_len, d_model)`.
+    x = self.dropout(x, training=training)
+    for block in self.enc_layers:
+      x = block([x, mask], training)
+    # x Shape `(batch_size, seq_len, d_model)`.
+    return x if self.output_all else x[:, 0, :]
+
+
+class TextEncoder(Layer):
+
+  def __init__(self, params, name='text_encoder', reuse=None, **kwargs):
+    super(TextEncoder, self).__init__(name=name, **kwargs)
+    self.separator = params.get_or_default('separator', ' ')
+    self.cls_token = '[CLS]' + self.separator
+    self.sep_token = self.separator + '[SEP]' + self.separator
+    params.transformer.output_all_token_embeddings = False
+    trans_params = Parameter.make_from_pb(params.transformer)
+    vocab_file = params.get_or_default('vocab_file', None)
+    self.vocab = None
+    self.default_token_id = params.get_or_default('default_token_id', 0)
+    if vocab_file is not None:
+      self.vocab = tf.feature_column.categorical_column_with_vocabulary_file(
+          'tokens',
+          vocabulary_file=vocab_file,
+          default_value=self.default_token_id)
+      logging.info('vocab file of TextEncoder(%s) is %s', name, vocab_file)
+      trans_params.vocab_size = self.vocab.vocabulary_size
+    self.encoder = TransformerEncoder(trans_params, name='transformer')
+
+  def call(self, inputs, training=None, **kwargs):
+    if type(inputs) not in (tuple, list):
+      inputs = [inputs]
+    inputs = [tf.squeeze(text) for text in inputs]
+    batch_size = tf.shape(inputs[0])
+    cls = tf.fill(batch_size, self.cls_token)
+    sep = tf.fill(batch_size, self.sep_token)
+    sentences = [cls]
+    for sentence in inputs:
+      sentences.append(sentence)
+      sentences.append(sep)
+    text = tf.strings.join(sentences)
+    tokens = tf.strings.split(text, self.separator)
+    if self.vocab is not None:
+      features = {'tokens': tokens}
+      token_ids = self.vocab._transform_feature(features)
+      token_ids = tf.sparse.to_dense(
+          token_ids, default_value=self.default_token_id, name='token_ids')
+      length = tf.shape(token_ids)[-1]
+      token_ids = tf.cond(
+          tf.less_equal(length, self.encoder.max_position), lambda: token_ids,
+          lambda: tf.slice(token_ids, [0, 0], [-1, self.encoder.max_position]))
+      mask = tf.not_equal(token_ids, self.default_token_id, name='mask')
+    else:
+      tokens = tf.sparse.to_dense(tokens, default_value='')
+      length = tf.shape(tokens)[-1]
+      tokens = tf.cond(
+          tf.less_equal(length, self.encoder.max_position), lambda: tokens,
+          lambda: tf.slice(tokens, [0, 0], [-1, self.encoder.max_position]))
+      token_ids = tf.string_to_hash_bucket_fast(
+          tokens, self.encoder.vocab_size, name='token_ids')
+      mask = tf.not_equal(tokens, '', name='mask')
+
+    encoding = self.encoder([token_ids, mask], training=training)
+    return encoding

easy_rec/python/layers/layer_norm.py

@@ -0,0 +1,51 @@
+# -*- encoding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class LayerNormalization(tf.layers.Layer):
+  """Layer normalization for BTC format: supports L2(default) and L1 modes."""
+
+  def __init__(self, hidden_size, params={}):
+    super(LayerNormalization, self).__init__()
+    self.hidden_size = hidden_size
+    self.norm_type = params.get('type', 'layernorm_L2')
+    self.epsilon = params.get('epsilon', 1e-6)
+
+  def build(self, _):
+    self.scale = tf.get_variable(
+        'layer_norm_scale', [self.hidden_size],
+        initializer=tf.keras.initializers.Ones(),
+        dtype=tf.float32)
+    self.bias = tf.get_variable(
+        'layer_norm_bias', [self.hidden_size],
+        initializer=tf.keras.initializers.Zeros(),
+        dtype=tf.float32)
+    self.built = True
+
+  def call(self, x):
+    if self.norm_type == 'layernorm_L2':
+      epsilon = self.epsilon
+      dtype = x.dtype
+      x = tf.cast(x=x, dtype=tf.float32)
+      mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
+      variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
+      norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
+      result = norm_x * self.scale + self.bias
+      return tf.cast(x=result, dtype=dtype)
+
+    else:
+      dtype = x.dtype
+      if dtype == tf.float16:
+        x = tf.cast(x, dtype=tf.float32)
+      mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
+      x = x - mean
+      variance = tf.reduce_mean(tf.abs(x), axis=[-1], keepdims=True)
+      norm_x = tf.div(x, variance + self.epsilon)
+      y = norm_x * self.scale + self.bias
+      if dtype == tf.float16:
+        y = tf.saturate_cast(y, dtype)
+      return y

easy_rec/python/layers/mmoe.py

@@ -0,0 +1,83 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+
+from easy_rec.python.layers import dnn
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class MMOE:
+
+  def __init__(self,
+               expert_dnn_config,
+               l2_reg,
+               num_task,
+               num_expert=None,
+               name='mmoe',
+               is_training=False):
+    """Initializes a `DNN` Layer.
+
+    Args:
+      expert_dnn_config: a instance or a list of easy_rec.python.protos.dnn_pb2.DNN,
+        if it is a list of configs, the param `num_expert` will be ignored,
+        if it is a single config, the number of experts will be specified by num_expert.
+      l2_reg: l2 regularizer.
+      num_task: number of tasks
+      num_expert: number of experts, default is the list length of expert_dnn_configs
+      name: scope of the DNN, so that the parameters could be separated from other dnns
+      is_training: train phase or not, impact batchnorm and dropout
+    """
+    if isinstance(expert_dnn_config, list):
+      self._expert_dnn_configs = expert_dnn_config
+      self._num_expert = len(expert_dnn_config)
+    else:
+      assert num_expert is not None and num_expert > 0, \
+          'param `num_expert` must be large than zero, when expert_dnn_config is not a list'
+      self._expert_dnn_configs = [expert_dnn_config] * num_expert
+      self._num_expert = num_expert
+    logging.info('num_expert: {0}'.format(self._num_expert))
+
+    self._num_task = num_task
+    self._l2_reg = l2_reg
+    self._name = name
+    self._is_training = is_training
+
+  @property
+  def num_expert(self):
+    return self._num_expert
+
+  def gate(self, unit, deep_fea, name):
+    fea = tf.layers.dense(
+        inputs=deep_fea,
+        units=unit,
+        kernel_regularizer=self._l2_reg,
+        name='%s/dnn' % name)
+    fea = tf.nn.softmax(fea, axis=1)
+    return fea
+
+  def __call__(self, deep_fea):
+    expert_fea_list = []
+    for expert_id in range(self._num_expert):
+      expert_dnn_config = self._expert_dnn_configs[expert_id]
+      expert_dnn = dnn.DNN(
+          expert_dnn_config,
+          self._l2_reg,
+          name='%s/expert_%d' % (self._name, expert_id),
+          is_training=self._is_training)
+      expert_fea = expert_dnn(deep_fea)
+      expert_fea_list.append(expert_fea)
+    experts_fea = tf.stack(expert_fea_list, axis=1)
+
+    task_input_list = []
+    for task_id in range(self._num_task):
+      gate = self.gate(
+          self._num_expert, deep_fea, name='%s/gate_%d' % (self._name, task_id))
+      gate = tf.expand_dims(gate, -1)
+      task_input = tf.multiply(experts_fea, gate)
+      task_input = tf.reduce_sum(task_input, axis=1)
+      task_input_list.append(task_input)
+    return task_input_list

easy_rec/python/layers/multihead_attention.py

@@ -0,0 +1,162 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class MultiHeadAttention:
+
+  def __init__(self, head_num, head_size, l2_reg, use_res=False, name=''):
+    """Initializes a `MultiHeadAttention` Layer.
+
+    Args:
+      head_num: The number of heads
+      head_size: The dimension of a head
+      l2_reg: l2 regularizer
+      use_res: Whether to use residual connections before output.
+      name: scope of the MultiHeadAttention, so that the parameters could be separated from other MultiHeadAttention
+    """
+    self._head_num = head_num
+    self._head_size = head_size
+    self._l2_reg = l2_reg
+    self._use_res = use_res
+    self._name = name
+
+  def _split_multihead_qkv(self, q, k, v):
+    """Split multiple heads.
+
+    Args:
+      q: Query matrix of shape [bs, feature_num, head_num * head_size].
+      k: Key matrix of shape [bs, feature_num, head_num * head_size].
+      v: Value matrix of shape [bs, feature_num, head_num * head_size].
+
+    Returns:
+      q: Query matrix of shape [bs, head_num, feature_num, head_size].
+      k: Key matrix of shape [bs, head_num, feature_num, head_size].
+      v: Value matrix of shape [bs, head_num, feature_num, head_size].
+    """
+    reshaped_q = tf.reshape(
+        q, shape=[-1, q.shape[1], self._head_num, self._head_size])
+    q = tf.transpose(reshaped_q, perm=[0, 2, 1, 3])
+    reshaped_k = tf.reshape(
+        k, shape=[-1, k.shape[1], self._head_num, self._head_size])
+    k = tf.transpose(reshaped_k, perm=[0, 2, 1, 3])
+    reshaped_v = tf.reshape(
+        v, shape=[-1, v.shape[1], self._head_num, self._head_size])
+    v = tf.transpose(reshaped_v, perm=[0, 2, 1, 3])
+    return q, k, v
+
+  def _scaled_dot_product_attention(self, q, k, v):
+    """Calculate scaled dot product attention by q, k and v.
+
+    Args:
+      q: Query matrix of shape [bs, head_num, feature_num, head_size].
+      k: Key matrix of shape [bs, head_num, feature_num, head_size].
+      v: Value matrix of shape [bs, head_num, feature_num, head_size].
+
+    Returns:
+      q: Query matrix of shape [bs, head_num, feature_num, head_size].
+      k: Key matrix of shape [bs, head_num, feature_num, head_size].
+      v: Value matrix of shape [bs, head_num, feature_num, head_size].
+    """
+    product = tf.linalg.matmul(
+        a=q, b=k, transpose_b=True) / (
+            self._head_size**-0.5)
+    weights = tf.nn.softmax(product)
+    out = tf.linalg.matmul(weights, v)
+    return out
+
+  def _compute_qkv(self, q, k, v):
+    """Calculate q, k and v matrices.
+
+    Args:
+      q: Query matrix of shape [bs, feature_num, d_model].
+      k: Key matrix of shape [bs, feature_num, d_model].
+      v: Value matrix of shape [bs, feature_num, d_model].
+
+    Returns:
+      q: Query matrix of shape [bs, feature_num, head_size * n_head].
+      k: Key matrix of shape [bs, feature_num, head_size * n_head].
+      v: Value matrix of shape [bs, feature_num, head_size * n_head].
+    """
+    q = tf.layers.dense(
+        q,
+        self._head_num * self._head_size,
+        use_bias=False,
+        kernel_regularizer=self._l2_reg,
+        name='%s/%s/dnn' % (self._name, 'query'))
+    k = tf.layers.dense(
+        k,
+        self._head_num * self._head_size,
+        use_bias=False,
+        kernel_regularizer=self._l2_reg,
+        name='%s/%s/dnn' % (self._name, 'key'))
+    v = tf.layers.dense(
+        v,
+        self._head_num * self._head_size,
+        use_bias=False,
+        kernel_regularizer=self._l2_reg,
+        name='%s/%s/dnn' % (self._name, 'value'))
+    return q, k, v
+
+  def _combine_heads(self, multi_head_tensor):
+    """Combine the results of multiple heads.
+
+    Args:
+      multi_head_tensor: Result matrix of shape [bs, head_num, feature_num, head_size].
+
+    Returns:
+      out: Result matrix of shape [bs, feature_num, head_num * head_size].
+    """
+    x = tf.transpose(multi_head_tensor, perm=[0, 2, 1, 3])
+    out = tf.reshape(x, shape=[-1, x.shape[1], x.shape[2] * x.shape[3]])
+    return out
+
+  def _multi_head_attention(self, attention_input):
+    """Build multiple heads attention layer.
+
+    Args:
+      attention_input: The input of interacting layer, has a shape of [bs, feature_num, d_model].
+
+    Returns:
+      out: The output of multi head attention layer, has a shape of [bs, feature_num, head_num * head_size].
+    """
+    if isinstance(attention_input, list):
+      assert len(attention_input) == 3 or len(attention_input) == 1, \
+          'If the input of multi_head_attention is a list, the length must be 1 or 3.'
+
+      if len(attention_input) == 3:
+        ori_q = attention_input[0]
+        ori_k = attention_input[1]
+        ori_v = attention_input[2]
+      else:
+        ori_q = attention_input[0]
+        ori_k = attention_input[0]
+        ori_v = attention_input[0]
+    else:
+      ori_q = attention_input
+      ori_k = attention_input
+      ori_v = attention_input
+
+    q, k, v = self._compute_qkv(ori_q, ori_k, ori_v)
+    q, k, v = self._split_multihead_qkv(q, k, v)
+    multi_head_tensor = self._scaled_dot_product_attention(q, k, v)
+    out = self._combine_heads(multi_head_tensor)
+
+    if self._use_res:
+      W_0_x = tf.layers.dense(
+          ori_v,
+          out.shape[2],
+          use_bias=False,
+          kernel_regularizer=self._l2_reg,
+          name='%s/dnn' % (self._name))
+      res_out = tf.nn.relu(out + W_0_x)
+      return res_out
+    else:
+      return out
+
+  def __call__(self, deep_fea):
+    deep_fea = self._multi_head_attention(deep_fea)
+    return deep_fea