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

easy_rec/python/layers/keras/bst.py

@@ -0,0 +1,119 @@
+# -*- encoding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.layers import multihead_cross_attention
+from easy_rec.python.utils.activation import get_activation
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class BST(Layer):
+
+  def __init__(self, params, name='bst', reuse=None, **kwargs):
+    super(BST, self).__init__(name=name, **kwargs)
+    self.reuse = reuse
+    self.l2_reg = params.l2_regularizer
+    self.config = params.get_pb_config()
+
+  def encode(self, seq_input, max_position):
+    seq_fea = multihead_cross_attention.embedding_postprocessor(
+        seq_input,
+        position_embedding_name=self.name,
+        max_position_embeddings=max_position,
+        reuse_position_embedding=self.reuse)
+
+    n = tf.count_nonzero(seq_input, axis=-1)
+    seq_mask = tf.cast(n > 0, tf.int32)
+
+    attention_mask = multihead_cross_attention.create_attention_mask_from_input_mask(
+        from_tensor=seq_fea, to_mask=seq_mask)
+
+    hidden_act = get_activation(self.config.hidden_act)
+    attention_fea = multihead_cross_attention.transformer_encoder(
+        seq_fea,
+        hidden_size=self.config.hidden_size,
+        num_hidden_layers=self.config.num_hidden_layers,
+        num_attention_heads=self.config.num_attention_heads,
+        attention_mask=attention_mask,
+        intermediate_size=self.config.intermediate_size,
+        intermediate_act_fn=hidden_act,
+        hidden_dropout_prob=self.config.hidden_dropout_prob,
+        attention_probs_dropout_prob=self.config.attention_probs_dropout_prob,
+        initializer_range=self.config.initializer_range,
+        name=self.name + '/transformer',
+        reuse=self.reuse)
+    # attention_fea shape: [batch_size, seq_length, hidden_size]
+    if self.config.output_all_token_embeddings:
+      out_fea = tf.reshape(attention_fea,
+                           [-1, max_position * self.config.hidden_size])
+    else:
+      out_fea = attention_fea[:, 0, :]  # target feature
+    print('bst output shape:', out_fea.shape)
+    return out_fea
+
+  def call(self, inputs, training=None, **kwargs):
+    if not training:
+      self.config.hidden_dropout_prob = 0.0
+      self.config.attention_probs_dropout_prob = 0.0
+    assert isinstance(inputs, (list, tuple))
+    assert len(inputs) >= 2
+    # seq_input: [batch_size, seq_len, embed_size]
+    seq_input, seq_len = inputs[:2]
+    target = inputs[2] if len(inputs) > 2 else None
+    max_position = self.config.max_position_embeddings
+    # max_seq_len: the max sequence length in current mini-batch, all sequences are padded to this length
+    batch_size, cur_batch_max_seq_len, seq_embed_size = get_shape_list(
+        seq_input, 3)
+    valid_len = tf.assert_less_equal(
+        cur_batch_max_seq_len,
+        max_position,
+        message='sequence length is greater than `max_position_embeddings`:' +
+        str(max_position) + ' in feature group:' + self.name +
+        ', you should set `max_seq_len` in sequence feature configs')
+
+    if self.config.output_all_token_embeddings:
+      seq_input = tf.cond(
+          tf.constant(max_position) > cur_batch_max_seq_len, lambda: tf.pad(
+              seq_input, [[0, 0], [0, max_position - cur_batch_max_seq_len],
+                          [0, 0]], 'CONSTANT'),
+          lambda: tf.slice(seq_input, [0, 0, 0], [-1, max_position, -1]))
+
+    if seq_embed_size != self.config.hidden_size:
+      seq_input = tf.layers.dense(
+          seq_input,
+          self.config.hidden_size,
+          activation=tf.nn.relu,
+          kernel_regularizer=self.l2_reg,
+          name=self.name + '/seq_project',
+          reuse=self.reuse)
+
+    keep_target = self.config.target_item_position in ('head', 'tail')
+    if target is not None and keep_target:
+      target_size = target.shape.as_list()[-1]
+      assert seq_embed_size == target_size, 'the embedding size of sequence and target item is not equal' \
+                                            ' in feature group:' + self.name
+      if target_size != self.config.hidden_size:
+        target = tf.layers.dense(
+            target,
+            self.config.hidden_size,
+            activation=tf.nn.relu,
+            kernel_regularizer=self.l2_reg,
+            name=self.name + '/target_project',
+            reuse=self.reuse)
+      # target_feature: [batch_size, 1, embed_size]
+      target = tf.expand_dims(target, 1)
+      # seq_input: [batch_size, seq_len+1, embed_size]
+      if self.config.target_item_position == 'head':
+        seq_input = tf.concat([target, seq_input], axis=1)
+      else:
+        seq_input = tf.concat([seq_input, target], axis=1)
+      max_position += 1
+    elif self.config.reserve_target_position:
+      max_position += 1
+
+    with tf.control_dependencies([valid_len]):
+      return self.encode(seq_input, max_position)

easy_rec/python/layers/keras/custom_ops.py

@@ -0,0 +1,250 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Convenience blocks for using custom ops."""
+import logging
+import os
+
+import tensorflow as tf
+from tensorflow.python.framework import ops
+from tensorflow.python.keras.layers import Layer
+
+curr_dir, _ = os.path.split(__file__)
+parent_dir = os.path.dirname(curr_dir)
+ops_idr = os.path.dirname(parent_dir)
+ops_dir = os.path.join(ops_idr, 'ops')
+if 'PAI' in tf.__version__:
+  ops_dir = os.path.join(ops_dir, '1.12_pai')
+elif tf.__version__.startswith('1.12'):
+  ops_dir = os.path.join(ops_dir, '1.12')
+elif tf.__version__.startswith('1.15'):
+  if 'IS_ON_PAI' in os.environ:
+    ops_dir = os.path.join(ops_dir, 'DeepRec')
+  else:
+    ops_dir = os.path.join(ops_dir, '1.15')
+elif tf.__version__.startswith('2.12'):
+  ops_dir = os.path.join(ops_dir, '2.12')
+
+logging.info('ops_dir is %s' % ops_dir)
+custom_op_path = os.path.join(ops_dir, 'libcustom_ops.so')
+try:
+  custom_ops = tf.load_op_library(custom_op_path)
+  logging.info('load custom op from %s succeed' % custom_op_path)
+except Exception as ex:
+  logging.warning('load custom op from %s failed: %s' %
+                  (custom_op_path, str(ex)))
+  custom_ops = None
+
+# if tf.__version__ >= '2.0':
+#   tf = tf.compat.v1
+
+
+class SeqAugmentOps(Layer):
+  """Do data augmentation for input sequence embedding."""
+
+  def __init__(self, params, name='sequence_aug', reuse=None, **kwargs):
+    super(SeqAugmentOps, self).__init__(name=name, **kwargs)
+    self.reuse = reuse
+    self.seq_aug_params = params.get_pb_config()
+    self.seq_augment = custom_ops.my_seq_augment
+
+  def call(self, inputs, training=None, **kwargs):
+    assert isinstance(
+        inputs,
+        (list, tuple)), 'the inputs of SeqAugmentOps must be type of list/tuple'
+    assert len(inputs) >= 2, 'SeqAugmentOps must have at least 2 inputs'
+    seq_input, seq_len = inputs[:2]
+    embedding_dim = int(seq_input.shape[-1])
+    with tf.variable_scope(self.name, reuse=self.reuse):
+      mask_emb = tf.get_variable(
+          'mask', (embedding_dim,), dtype=tf.float32, trainable=True)
+    seq_len = tf.to_int32(seq_len)
+    with ops.device('/CPU:0'):
+      aug_seq, aug_len = self.seq_augment(seq_input, seq_len, mask_emb,
+                                          self.seq_aug_params.crop_rate,
+                                          self.seq_aug_params.reorder_rate,
+                                          self.seq_aug_params.mask_rate)
+    return aug_seq, aug_len
+
+
+class TextNormalize(Layer):
+
+  def __init__(self, params, name='text_normalize', reuse=None, **kwargs):
+    super(TextNormalize, self).__init__(name=name, **kwargs)
+    self.txt_normalizer = custom_ops.text_normalize_op
+    self.norm_parameter = params.get_or_default('norm_parameter', 0)
+    self.remove_space = params.get_or_default('remove_space', False)
+
+  def call(self, inputs, training=None, **kwargs):
+    inputs = inputs if type(inputs) in (tuple, list) else [inputs]
+    with ops.device('/CPU:0'):
+      result = [
+          self.txt_normalizer(
+              txt,
+              parameter=self.norm_parameter,
+              remove_space=self.remove_space) for txt in inputs
+      ]
+    if len(result) == 1:
+      return result[0]
+    return result
+
+
+class MappedDotProduct(Layer):
+
+  def __init__(self, params, name='mapped_dot_product', reuse=None, **kwargs):
+    super(MappedDotProduct, self).__init__(name=name, **kwargs)
+    self.mapped_dot_product = custom_ops.mapped_dot_product
+    self.bucketize = custom_ops.my_bucketize
+    self.default_value = params.get_or_default('default_value', 0)
+    self.separator = params.get_or_default('separator', '\035')
+    self.norm_fn = params.get_or_default('normalize_fn', None)
+    self.boundaries = list(params.get_or_default('boundaries', []))
+    self.emb_dim = params.get_or_default('embedding_dim', 0)
+    self.print_first_n = params.get_or_default('print_first_n', 0)
+    self.summarize = params.get_or_default('summarize', None)
+    if self.emb_dim > 0:
+      vocab_size = len(self.boundaries) + 1
+      with tf.variable_scope(self.name, reuse=reuse):
+        self.embedding_table = tf.get_variable(
+            name='dot_product_emb_table',
+            shape=[vocab_size, self.emb_dim],
+            dtype=tf.float32)
+
+  def call(self, inputs, training=None, **kwargs):
+    query, doc = inputs[:2]
+    with ops.device('/CPU:0'):
+      feature = self.mapped_dot_product(
+          query=query,
+          document=doc,
+          feature_name=self.name,
+          separator=self.separator,
+          default_value=self.default_value)
+      tf.summary.scalar(self.name, tf.reduce_mean(feature))
+      if self.print_first_n:
+        encode_q = tf.regex_replace(query, self.separator, ' ')
+        encode_t = tf.regex_replace(query, self.separator, ' ')
+        feature = tf.Print(
+            feature, [encode_q, encode_t, feature],
+            message=self.name,
+            first_n=self.print_first_n,
+            summarize=self.summarize)
+      if self.norm_fn is not None:
+        fn = eval(self.norm_fn)
+        feature = fn(feature)
+        tf.summary.scalar('normalized_%s' % self.name, tf.reduce_mean(feature))
+        if self.print_first_n:
+          feature = tf.Print(
+              feature, [feature],
+              message='normalized %s' % self.name,
+              first_n=self.print_first_n,
+              summarize=self.summarize)
+      if self.boundaries:
+        feature = self.bucketize(feature, boundaries=self.boundaries)
+        tf.summary.histogram('bucketized_%s' % self.name, feature)
+    if self.emb_dim > 0 and self.boundaries:
+      vocab_size = len(self.boundaries) + 1
+      one_hot_input_ids = tf.one_hot(feature, depth=vocab_size)
+      return tf.matmul(one_hot_input_ids, self.embedding_table)
+    return tf.expand_dims(feature, axis=-1)
+
+
+class OverlapFeature(Layer):
+
+  def __init__(self, params, name='overlap_feature', reuse=None, **kwargs):
+    super(OverlapFeature, self).__init__(name=name, **kwargs)
+    self.overlap_feature = custom_ops.overlap_fg_op
+    methods = params.get_or_default('methods', [])
+    assert methods, 'overlap feature methods must be set'
+    self.methods = [str(method) for method in methods]
+    self.norm_fn = params.get_or_default('normalize_fn', None)
+    self.boundaries = list(params.get_or_default('boundaries', []))
+    self.separator = params.get_or_default('separator', '\035')
+    self.default_value = params.get_or_default('default_value', '-1')
+    self.emb_dim = params.get_or_default('embedding_dim', 0)
+    self.print_first_n = params.get_or_default('print_first_n', 0)
+    self.summarize = params.get_or_default('summarize', None)
+    if self.emb_dim > 0:
+      vocab_size = len(self.boundaries) + 1
+      vocab_size *= len(self.methods)
+      with tf.variable_scope(self.name, reuse=reuse):
+        self.embedding_table = tf.get_variable(
+            name='overlap_emb_table',
+            shape=[vocab_size, self.emb_dim],
+            dtype=tf.float32)
+
+  def call(self, inputs, training=None, **kwargs):
+    query, title = inputs[:2]
+    with ops.device('/CPU:0'):
+      feature = self.overlap_feature(
+          query=query,
+          title=title,
+          feature_name=self.name,
+          separator=self.separator,
+          default_value=self.default_value,
+          boundaries=self.boundaries,
+          methods=self.methods,
+          dtype=tf.int32 if self.boundaries else tf.float32)
+
+    for i, method in enumerate(self.methods):
+      # warning: feature[:, i] may be not the result of method
+      if self.boundaries:
+        tf.summary.histogram('bucketized_%s' % method, feature[:, i])
+      else:
+        tf.summary.scalar(method, tf.reduce_mean(feature[:, i]))
+    if self.print_first_n:
+      encode_q = tf.regex_replace(query, self.separator, ' ')
+      encode_t = tf.regex_replace(query, self.separator, ' ')
+      feature = tf.Print(
+          feature, [encode_q, encode_t, feature],
+          message=self.name,
+          first_n=self.print_first_n,
+          summarize=self.summarize)
+    if self.norm_fn is not None:
+      fn = eval(self.norm_fn)
+      feature = fn(feature)
+
+    if self.emb_dim > 0 and self.boundaries:
+      # This vocab will be small so we always do one-hot here, since it is always
+      # faster for a small vocabulary.
+      batch_size = tf.shape(feature)[0]
+      vocab_size = len(self.boundaries) + 1
+      num_indices = len(self.methods)
+      # Compute offsets, add to every column indices
+      offsets = tf.range(num_indices) * vocab_size  # Shape: [3]
+      offsets = tf.reshape(offsets, [1, num_indices])  # Shape: [1, 3]
+      offsets = tf.tile(offsets,
+                        [batch_size, 1])  # Shape: [batch_size, num_indices]
+      shifted_indices = feature + offsets  # Shape: [batch_size, num_indices]
+      flat_feature_ids = tf.reshape(shifted_indices, [-1])
+      one_hot_ids = tf.one_hot(flat_feature_ids, depth=vocab_size * num_indices)
+      feature_embeddings = tf.matmul(one_hot_ids, self.embedding_table)
+      feature_embeddings = tf.reshape(feature_embeddings,
+                                      [batch_size, num_indices * self.emb_dim])
+      return feature_embeddings
+    return feature
+
+
+class EditDistance(Layer):
+
+  def __init__(self, params, name='edit_distance', reuse=None, **kwargs):
+    super(EditDistance, self).__init__(name=name, **kwargs)
+    self.edit_distance = custom_ops.my_edit_distance
+    self.txt_encoding = params.get_or_default('text_encoding', 'utf-8')
+    self.emb_size = params.get_or_default('embedding_size', 512)
+    emb_dim = params.get_or_default('embedding_dim', 4)
+    with tf.variable_scope(self.name, reuse=reuse):
+      self.embedding_table = tf.get_variable('embedding_table',
+                                             [self.emb_size, emb_dim],
+                                             tf.float32)
+
+  def call(self, inputs, training=None, **kwargs):
+    input1, input2 = inputs[:2]
+    with ops.device('/CPU:0'):
+      dist = self.edit_distance(
+          input1,
+          input2,
+          normalize=False,
+          dtype=tf.int32,
+          encoding=self.txt_encoding)
+    ids = tf.clip_by_value(dist, 0, self.emb_size - 1)
+    embed = tf.nn.embedding_lookup(self.embedding_table, ids)
+    return embed

easy_rec/python/layers/keras/data_augment.py

@@ -0,0 +1,133 @@
+# -*- encoding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+def item_mask(aug_data, length, mask_emb, mask_rate):
+  length1 = tf.cast(length, dtype=tf.float32)
+  num_mask = tf.cast(tf.math.floor(length1 * mask_rate), dtype=tf.int32)
+  max_len = tf.shape(aug_data)[0]
+  seq_mask = tf.sequence_mask(num_mask, length)
+  seq_mask = tf.random.shuffle(seq_mask)
+  padding = tf.sequence_mask(0, max_len - length)
+  seq_mask = tf.concat([seq_mask, padding], axis=0)
+
+  mask_emb = tf.tile(mask_emb, [max_len, 1])
+
+  masked_item_seq = tf.where(seq_mask, mask_emb, aug_data)
+  return masked_item_seq, length
+
+
+def item_crop(aug_data, length, crop_rate):
+  length1 = tf.cast(length, dtype=tf.float32)
+  max_len, _ = get_shape_list(aug_data)
+  max_length = tf.cast(max_len, dtype=tf.int32)
+
+  num_left = tf.cast(tf.math.floor(length1 * crop_rate), dtype=tf.int32)
+  crop_begin = tf.random.uniform([],
+                                 minval=0,
+                                 maxval=length - num_left,
+                                 dtype=tf.int32)
+  zeros = tf.zeros_like(aug_data)
+  x = aug_data[crop_begin:crop_begin + num_left]
+  y = zeros[:max_length - num_left]
+  cropped = tf.concat([x, y], axis=0)
+  cropped_item_seq = tf.where(
+      crop_begin + num_left < max_length, cropped,
+      tf.concat([aug_data[crop_begin:], zeros[:crop_begin]], axis=0))
+  return cropped_item_seq, num_left
+
+
+def item_reorder(aug_data, length, reorder_rate):
+  length1 = tf.cast(length, dtype=tf.float32)
+  num_reorder = tf.cast(tf.math.floor(length1 * reorder_rate), dtype=tf.int32)
+  reorder_begin = tf.random.uniform([],
+                                    minval=0,
+                                    maxval=length - num_reorder,
+                                    dtype=tf.int32)
+  shuffle_index = tf.range(reorder_begin, reorder_begin + num_reorder)
+  shuffle_index = tf.random.shuffle(shuffle_index)
+  x = tf.range(get_shape_list(aug_data)[0])
+  left = tf.slice(x, [0], [reorder_begin])
+  right = tf.slice(x, [reorder_begin + num_reorder], [-1])
+  reordered_item_index = tf.concat([left, shuffle_index, right], axis=0)
+  reordered_item_seq = tf.scatter_nd(
+      tf.expand_dims(reordered_item_index, axis=1), aug_data,
+      tf.shape(aug_data))
+  return reordered_item_seq, length
+
+
+def augment_fn(x, aug_param, mask):
+  seq, length = x
+
+  def crop_fn():
+    return item_crop(seq, length, aug_param.crop_rate)
+
+  def mask_fn():
+    return item_mask(seq, length, mask, aug_param.mask_rate)
+
+  def reorder_fn():
+    return item_reorder(seq, length, aug_param.reorder_rate)
+
+  trans_fn = []
+  if aug_param.crop_rate < 1.0:
+    trans_fn.append(crop_fn)
+  if aug_param.mask_rate > 0:
+    trans_fn.append(mask_fn)
+  if aug_param.reorder_rate > 0:
+    trans_fn.append(reorder_fn)
+
+  num_trans = len(trans_fn)
+  if num_trans == 0:
+    return seq, length
+
+  if num_trans == 1:
+    return trans_fn[0]()
+
+  method = tf.random.uniform([], minval=0, maxval=num_trans, dtype=tf.int32)
+  if num_trans == 2:
+    return tf.cond(tf.equal(method, 0), trans_fn[0], trans_fn[1])
+
+  aug_seq, aug_len = tf.cond(
+      tf.equal(method, 0), crop_fn,
+      lambda: tf.cond(tf.equal(method, 1), mask_fn, reorder_fn))
+  return aug_seq, aug_len
+
+
+def sequence_augment(seq_input, seq_len, mask, aug_param):
+  lengths = tf.cast(seq_len, dtype=tf.int32)
+  aug_seq, aug_len = tf.map_fn(
+      lambda elems: augment_fn(elems, aug_param, mask),
+      elems=(seq_input, lengths),
+      dtype=(tf.float32, tf.int32))
+
+  aug_seq = tf.reshape(aug_seq, tf.shape(seq_input))
+  return aug_seq, aug_len
+
+
+class SeqAugment(Layer):
+  """Do data augmentation for input sequence embedding."""
+
+  def __init__(self, params, name='seq_aug', reuse=None, **kwargs):
+    super(SeqAugment, self).__init__(name=name, **kwargs)
+    self.reuse = reuse
+    self.seq_aug_params = params.get_pb_config()
+
+  def call(self, inputs, training=None, **kwargs):
+    assert isinstance(inputs, (list, tuple))
+    seq_input, seq_len = inputs[:2]
+
+    embedding_size = int(seq_input.shape[-1])
+    with tf.variable_scope(self.name, reuse=self.reuse):
+      mask_emb = tf.get_variable(
+          'mask', [1, embedding_size], dtype=tf.float32, trainable=True)
+
+    aug_seq, aug_len = sequence_augment(seq_input, seq_len, mask_emb,
+                                        self.seq_aug_params)
+    return aug_seq, aug_len

easy_rec/python/layers/keras/din.py

@@ -0,0 +1,67 @@
+# -*- encoding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
+
+import tensorflow as tf
+from tensorflow.python.keras.layers import Layer
+
+from easy_rec.python.layers.keras import MLP
+from easy_rec.python.layers.utils import Parameter
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+
+class DIN(Layer):
+
+  def __init__(self, params, name='din', reuse=None, **kwargs):
+    super(DIN, self).__init__(name=name, **kwargs)
+    self.reuse = reuse
+    self.l2_reg = params.l2_regularizer
+    self.config = params.get_pb_config()
+    self.config.attention_dnn.use_final_bn = False
+    self.config.attention_dnn.use_final_bias = True
+    self.config.attention_dnn.final_activation = 'linear'
+    mlp_params = Parameter.make_from_pb(self.config.attention_dnn)
+    mlp_params.l2_regularizer = self.l2_reg
+    self.din_layer = MLP(mlp_params, 'din_attention', reuse=self.reuse)
+
+  def call(self, inputs, training=None, **kwargs):
+    keys, seq_len, query = inputs
+    assert query is not None, '[%s] target feature is empty' % self.name
+    query_emb_size = int(query.shape[-1])
+    seq_emb_size = keys.shape.as_list()[-1]
+    if query_emb_size != seq_emb_size:
+      logging.info(
+          '<din> the embedding size of sequence [%d] and target item [%d] is not equal'
+          ' in feature group: %s', seq_emb_size, query_emb_size, self.name)
+      if query_emb_size < seq_emb_size:
+        query = tf.pad(query, [[0, 0], [0, seq_emb_size - query_emb_size]])
+      else:
+        assert False, 'the embedding size of target item is larger than the one of sequence'
+
+    batch_size, max_seq_len, _ = get_shape_list(keys, 3)
+    queries = tf.tile(tf.expand_dims(query, 1), [1, max_seq_len, 1])
+    din_all = tf.concat([queries, keys, queries - keys, queries * keys],
+                        axis=-1)
+    output = self.din_layer(din_all, training)  # [B, L, 1]
+    scores = tf.transpose(output, [0, 2, 1])  # [B, 1, L]
+
+    seq_mask = tf.sequence_mask(seq_len, max_seq_len, dtype=tf.bool)
+    seq_mask = tf.expand_dims(seq_mask, 1)
+    paddings = tf.ones_like(scores) * (-2**32 + 1)
+    scores = tf.where(seq_mask, scores, paddings)  # [B, 1, L]
+    if self.config.attention_normalizer == 'softmax':
+      scores = tf.nn.softmax(scores)  # (B, 1, L)
+    elif self.config.attention_normalizer == 'sigmoid':
+      scores = scores / (seq_emb_size**0.5)
+      scores = tf.nn.sigmoid(scores)
+    else:
+      raise ValueError('unsupported attention normalizer: ' +
+                       self.config.attention_normalizer)
+
+    if query_emb_size < seq_emb_size:
+      keys = keys[:, :, :query_emb_size]  # [B, L, E]
+    output = tf.squeeze(tf.matmul(scores, keys), axis=[1])
+    if self.config.need_target_feature:
+      output = tf.concat([output, query], axis=-1)
+    print('din output shape:', output.shape)
+    return output