models.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# author： JMXGODLZZ
# datetime： 2022/3/25 上午11:50 
# ide： PyCharm
# ! -*- coding: utf-8 -*-
# 自定义层

import numpy as np
import tensorflow as tf
from bert4keras.backend import keras, K, is_tf_keras
from bert4keras.backend import sequence_masking
from bert4keras.backend import recompute_grad
from keras import initializers, activations
from keras.layers import *


def integerize_shape(func):
    """装饰器，保证input_shape一定是int或None
    """

    def convert(item):
        if hasattr(item, '__iter__'):
            return [convert(i) for i in item]
        elif hasattr(item, 'value'):
            return item.value
        else:
            return item

    def new_func(self, input_shape):
        input_shape = convert(input_shape)
        return func(self, input_shape)

    return new_func


if (not is_tf_keras) and keras.__version__ < '2.3':

    class Layer(keras.layers.Layer):
        """重新定义Layer，赋予“层中层”功能
        （仅keras 2.3以下版本需要）
        """

        def __init__(self, **kwargs):
            super(Layer, self).__init__(**kwargs)
            self.supports_masking = True  # 本项目的自定义层均可mask

        def __setattr__(self, name, value):
            if isinstance(value, keras.layers.Layer):
                if not hasattr(self, '_layers'):
                    self._layers = []
                if value not in self._layers:
                    self._layers.append(value)
            super(Layer, self).__setattr__(name, value)

        @property
        def trainable_weights(self):
            trainable = getattr(self, 'trainable', True)
            if trainable:
                trainable_weights = super(Layer, self).trainable_weights[:]
                for l in getattr(self, '_layers', []):
                    trainable_weights += l.trainable_weights
                return trainable_weights
            else:
                return []

        @property
        def non_trainable_weights(self):
            trainable = getattr(self, 'trainable', True)
            non_trainable_weights = super(Layer, self).non_trainable_weights[:]
            for l in getattr(self, '_layers', []):
                if trainable:
                    non_trainable_weights += l.non_trainable_weights
                else:
                    non_trainable_weights += l.weights
            return non_trainable_weights


    if keras.__version__ < '2.2.5':

        import inspect


        class Model(keras.models.Model):
            """重新定义Model，整合fit和fit_generator
            """

            def fit(self, x=None, *args, **kwargs):
                if inspect.isgenerator(x):
                    return self.fit_generator(x, *args, **kwargs)
                else:
                    return super(Model, self).fit(x, *args, **kwargs)


        keras.models.Model = Model

else:

    class Layer(keras.layers.Layer):
        def __init__(self, **kwargs):
            super(Layer, self).__init__(**kwargs)
            self.supports_masking = True  # 本项目的自定义层均可mask

if (not is_tf_keras) or tf.__version__ < '1.15':

    if not is_tf_keras:
        NodeBase = keras.engine.base_layer.Node
    else:
        from tensorflow.python.keras.engine import base_layer

        NodeBase = base_layer.Node


    class Node(NodeBase):
        """修改Node来修复keras下孪生网络的bug
        注意：这是keras的bug，并不是bert4keras的bug，但keras已经不更新了，
              所以只好在这里进行修改。tf 1.15+自带的keras已经修改了这个
              bug。
        """

        @property
        def arguments(self):
            return self._arguments.copy()

        @arguments.setter
        def arguments(self, value):
            self._arguments = value or {}


    if not is_tf_keras:
        keras.engine.base_layer.Node = Node
    else:
        base_layer.Node = Node


class GlobalAveragePooling1D(keras.layers.GlobalAveragePooling1D):
    """重新定义GlobalAveragePooling1D，支持序列长度为None
    """

    def call(self, inputs, mask=None):
        axis = 1 if self.data_format == 'channels_last' else 2
        if mask is not None:
            mask = K.cast(mask, K.floatx())
            mask = mask[..., None] if axis == 1 else mask[:, None]
            return K.sum(inputs * mask, axis=axis) / K.sum(mask, axis=axis)
        else:
            return K.mean(inputs, axis=axis)


class GlobalMaxPooling1D(keras.layers.GlobalMaxPooling1D):
    """重新定义GlobalMaxPooling1D，支持mask
    """

    def __init__(self, data_format='channels_last', **kwargs):
        super(GlobalMaxPooling1D, self).__init__(data_format, **kwargs)
        self.supports_masking = True

    def call(self, inputs, mask=None):
        axis = 1 if self.data_format == 'channels_last' else 2
        inputs = sequence_masking(inputs, mask, '-inf', axis)
        return K.max(inputs, axis=axis)

    def compute_mask(self, inputs, mask=None):
        return None


# 直接覆盖原对象
keras.layers.GlobalAveragePooling1D = GlobalAveragePooling1D
keras.layers.GlobalMaxPooling1D = GlobalMaxPooling1D


class Embedding(keras.layers.Embedding):
    """拓展Embedding层
    """

    def compute_mask(self, inputs, mask=None):
        """为了适配T5，保证第一个token不被mask
        """
        if K.ndim(inputs) == 2:
            mask = super(Embedding, self).compute_mask(inputs, mask)
            if mask is not None:
                mask1 = K.ones_like(mask[:, :1], dtype='bool')
                mask2 = mask[:, 1:]
                return K.concatenate([mask1, mask2], 1)
        else:
            return mask

    def call(self, inputs, mode='embedding'):
        """新增mode参数，可以为embedding或dense。如果为embedding，
        则等价于普通Embedding层；如果为dense，则等价于无bias的Dense层。
        """
        if mode == 'embedding':
            return super(Embedding, self).call(inputs)
        else:
            kernel = K.transpose(self.embeddings)
            return K.dot(inputs, kernel)

    def compute_output_shape(self, input_shape):
        """关于判据，本来是通过缓存call时的mode参数来判断的，但是后来发现
        Keras在使用compute_output_shape的时候不一定配套调用了call函数，
        所以缓存的mode可能是不准的，因此只能出此下策。
        """
        if len(input_shape) == 2:
            return super(Embedding, self).compute_output_shape(input_shape)
        else:
            return input_shape[:2] + (K.int_shape(self.embeddings)[0],)


class BiasAdd(Layer):
    """加上偏置项
    """

    @integerize_shape
    def build(self, input_shape):
        super(BiasAdd, self).build(input_shape)
        output_dim = input_shape[-1]
        self.bias = self.add_weight(
            name='bias',
            shape=(output_dim,),
            initializer='zeros',
            trainable=True
        )

    def call(self, inputs):
        return K.bias_add(inputs, self.bias)


class Concatenate1D(Layer):
    """1维序列拼接层
    说明：本来该功能可以直接通过Concatenate层来实现，无奈Keras
          自带的Concatenate层的compute_mask写得不合理，导致一个
          带mask的序列与一个不带mask的序列拼接会报错，因此干脆
          自己重写一个好了。
    """

    def call(self, inputs):
        return K.concatenate(inputs, axis=1)

    def compute_mask(self, inputs, mask=None):
        if mask is not None:
            masks = []
            for i, m in enumerate(mask):
                if m is None:
                    m = K.ones_like(inputs[i][..., 0], dtype='bool')
                masks.append(m)
            return K.concatenate(masks, axis=1)

    def compute_output_shape(self, input_shape):
        if all([shape[1] for shape in input_shape]):
            seq_len = sum([shape[1] for shape in input_shape])
            return (input_shape[0][0], seq_len, input_shape[0][2])
        else:
            return (input_shape[0][0], None, input_shape[0][2])


class MultiHeadAttention(Layer):
    """多头注意力机制
    """

    def __init__(
            self,
            heads,
            head_size,
            out_dim=None,
            key_size=None,
            use_bias=True,
            attention_scale=True,
            return_attention_scores=False,
            kernel_initializer='glorot_uniform',
            **kwargs
    ):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.out_dim = out_dim or heads * head_size
        self.key_size = key_size or head_size
        self.use_bias = use_bias
        self.attention_scale = attention_scale
        self.return_attention_scores = return_attention_scores
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        super(MultiHeadAttention, self).build(input_shape)
        self.q_dense = Dense(
            units=self.key_size * self.heads,
            use_bias=self.use_bias,
            kernel_initializer=self.kernel_initializer
        )
        self.k_dense = Dense(
            units=self.key_size * self.heads,
            use_bias=self.use_bias,
            kernel_initializer=self.kernel_initializer
        )
        self.v_dense = Dense(
            units=self.head_size * self.heads,
            use_bias=self.use_bias,
            kernel_initializer=self.kernel_initializer
        )
        self.o_dense = Dense(
            units=self.out_dim,
            use_bias=self.use_bias,
            kernel_initializer=self.kernel_initializer
        )

    @recompute_grad
    def call(self, inputs, mask=None, **kwargs):
        """实现多头注意力
        q_mask: 对输入的query序列的mask。
                主要是将输出结果的padding部分置0。
        v_mask: 对输入的value序列的mask。
                主要是防止attention读取到padding信息。
        """
        q, k, v = inputs[:3]
        q_mask, v_mask = None, None
        if mask is not None:
            q_mask, v_mask = mask[0], mask[2]
        # 线性变换
        qw = self.q_dense(q)
        kw = self.k_dense(k)
        vw = self.v_dense(v)
        # 形状变换
        qw = K.reshape(qw, (-1, K.shape(q)[1], self.heads, self.key_size))
        kw = K.reshape(kw, (-1, K.shape(k)[1], self.heads, self.key_size))
        vw = K.reshape(vw, (-1, K.shape(v)[1], self.heads, self.head_size))
        # Attention
        qkv_inputs = [qw, kw, vw] + inputs[3:]
        qv_masks = [q_mask, v_mask]
        o, a = self.pay_attention_to(qkv_inputs, qv_masks, **kwargs)
        # 完成输出
        o = K.reshape(o, (-1, K.shape(o)[1], self.head_size * self.heads))
        o = self.o_dense(o)
        # 返回结果
        if self.return_attention_scores:
            return [o, a]
        else:
            return o

    def pay_attention_to(self, inputs, mask=None, **kwargs):
        """实现标准的乘性多头注意力
        a_bias: 对attention矩阵的bias。
                不同的attention bias对应不同的应用。
        p_bias: 在attention里的位置偏置。
                一般用来指定相对位置编码的种类。
        说明: 这里单独分离出pay_attention_to函数，是为了方便
              继承此类来定义不同形式的atttention；此处要求
              返回o.shape=(batch_size, seq_len, heads, head_size)。
        """
        (qw, kw, vw), n = inputs[:3], 3
        q_mask, v_mask = mask
        a_bias, p_bias = kwargs.get('a_bias'), kwargs.get('p_bias')
        if a_bias:
            a_bias = inputs[n]
            n += 1
        if p_bias == 'rotary':
            cos_pos = K.repeat_elements(inputs[n][..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(inputs[n][..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # Attention
        a = tf.einsum('bjhd,bkhd->bhjk', qw, kw)
        # 处理位置编码
        if p_bias == 'typical_relative':
            position_bias = inputs[n]
            a = a + tf.einsum('bjhd,jkd->bhjk', qw, position_bias)
        elif p_bias == 't5_relative':
            position_bias = K.permute_dimensions(inputs[n], (2, 0, 1))
            a = a + K.expand_dims(position_bias, 0)
        # Attention（续）
        if self.attention_scale:
            a = a / self.key_size ** 0.5
        if a_bias is not None:
            a = a + a_bias
        a = sequence_masking(a, v_mask, '-inf', -1)
        A = K.softmax(a)
        # 完成输出
        o = tf.einsum('bhjk,bkhd->bjhd', A, vw)
        if p_bias == 'typical_relative':
            o = o + tf.einsum('bhjk,jkd->bjhd', A, position_bias)
        return o, a

    def compute_output_shape(self, input_shape):
        o_shape = (input_shape[0][0], input_shape[0][1], self.out_dim)
        if self.return_attention_scores:
            a_shape = (
                input_shape[0][0], self.heads, input_shape[0][1],
                input_shape[1][1]
            )
            return [o_shape, a_shape]
        else:
            return o_shape

    def compute_mask(self, inputs, mask=None):
        if mask is not None:
            if self.return_attention_scores:
                return [mask[0], None]
            else:
                return mask[0]

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'out_dim': self.out_dim,
            'key_size': self.key_size,
            'use_bias': self.use_bias,
            'attention_scale': self.attention_scale,
            'return_attention_scores': self.return_attention_scores,
            'kernel_initializer':
                initializers.serialize(self.kernel_initializer),
        }
        base_config = super(MultiHeadAttention, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class LayerNormalization(Layer):
    """(Conditional) Layer Normalization
    hidden_*系列参数仅为有条件输入时(conditional=True)使用
    """

    def __init__(
            self,
            center=True,
            scale=True,
            epsilon=None,
            conditional=False,
            hidden_units=None,
            hidden_activation='linear',
            hidden_initializer='glorot_uniform',
            **kwargs
    ):
        super(LayerNormalization, self).__init__(**kwargs)
        self.center = center
        self.scale = scale
        self.conditional = conditional
        self.hidden_units = hidden_units
        self.hidden_activation = activations.get(hidden_activation)
        self.hidden_initializer = initializers.get(hidden_initializer)
        self.epsilon = epsilon or 1e-12

    def compute_mask(self, inputs, mask=None):
        if self.conditional:
            masks = mask if mask is not None else []
            masks = [m[None] for m in masks if m is not None]
            if len(masks) == 0:
                return None
            else:
                return K.all(K.concatenate(masks, axis=0), axis=0)
        else:
            return mask

    def build(self, input_shape):
        super(LayerNormalization, self).build(input_shape)

        if self.conditional:
            shape = (input_shape[0][-1],)
        else:
            shape = (input_shape[-1],)

        if self.center:
            self.beta = self.add_weight(
                shape=shape, initializer='zeros', name='beta'
            )
        if self.scale:
            self.gamma = self.add_weight(
                shape=shape, initializer='ones', name='gamma'
            )

        if self.conditional:

            if self.hidden_units is not None:
                self.hidden_dense = Dense(
                    units=self.hidden_units,
                    activation=self.hidden_activation,
                    use_bias=False,
                    kernel_initializer=self.hidden_initializer
                )

            if self.center:
                self.beta_dense = Dense(
                    units=shape[0], use_bias=False, kernel_initializer='zeros'
                )
            if self.scale:
                self.gamma_dense = Dense(
                    units=shape[0], use_bias=False, kernel_initializer='zeros'
                )

    @recompute_grad
    def call(self, inputs):
        """如果是条件Layer Norm，则默认以list为输入，第二个是condition
        """
        if self.conditional:
            inputs, cond = inputs
            if self.hidden_units is not None:
                cond = self.hidden_dense(cond)
            for _ in range(K.ndim(inputs) - K.ndim(cond)):
                cond = K.expand_dims(cond, 1)
            if self.center:
                beta = self.beta_dense(cond) + self.beta
            if self.scale:
                gamma = self.gamma_dense(cond) + self.gamma
        else:
            if self.center:
                beta = self.beta
            if self.scale:
                gamma = self.gamma

        outputs = inputs
        if self.center:
            mean = K.mean(outputs, axis=-1, keepdims=True)
            outputs = outputs - mean
        if self.scale:
            variance = K.mean(K.square(outputs), axis=-1, keepdims=True)
            std = K.sqrt(variance + self.epsilon)
            outputs = outputs / std * gamma
        if self.center:
            outputs = outputs + beta

        return outputs

    def compute_output_shape(self, input_shape):
        if self.conditional:
            return input_shape[0]
        else:
            return input_shape

    def get_config(self):
        config = {
            'center': self.center,
            'scale': self.scale,
            'epsilon': self.epsilon,
            'conditional': self.conditional,
            'hidden_units': self.hidden_units,
            'hidden_activation': activations.serialize(self.hidden_activation),
            'hidden_initializer':
                initializers.serialize(self.hidden_initializer),
        }
        base_config = super(LayerNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class PositionEmbedding(Layer):
    """定义可训练的位置Embedding
    """

    def __init__(
            self,
            input_dim,
            output_dim,
            merge_mode='add',
            hierarchical=None,
            embeddings_initializer='zeros',
            custom_position_ids=False,
            **kwargs
    ):
        super(PositionEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.merge_mode = merge_mode
        self.hierarchical = hierarchical
        self.embeddings_initializer = initializers.get(embeddings_initializer)
        self.custom_position_ids = custom_position_ids

    def build(self, input_shape):
        super(PositionEmbedding, self).build(input_shape)
        self.embeddings = self.add_weight(
            name='embeddings',
            shape=(self.input_dim, self.output_dim),
            initializer=self.embeddings_initializer
        )

    def call(self, inputs):
        """如果custom_position_ids，那么第二个输入为自定义的位置id
        """
        if self.custom_position_ids:
            inputs, position_ids = inputs
            if 'int' not in K.dtype(position_ids):
                position_ids = K.cast(position_ids, 'int32')
        else:
            input_shape = K.shape(inputs)
            batch_size, seq_len = input_shape[0], input_shape[1]
            position_ids = K.arange(0, seq_len, dtype='int32')[None]

        if self.hierarchical:
            alpha = 0.4 if self.hierarchical is True else self.hierarchical
            embeddings = self.embeddings - alpha * self.embeddings[:1]
            embeddings = embeddings / (1 - alpha)
            embeddings_x = K.gather(embeddings, position_ids // self.input_dim)
            embeddings_y = K.gather(embeddings, position_ids % self.input_dim)
            embeddings = alpha * embeddings_x + (1 - alpha) * embeddings_y
        else:
            if self.custom_position_ids:
                embeddings = K.gather(self.embeddings, position_ids)
            else:
                embeddings = self.embeddings[None, :seq_len]

        if self.merge_mode == 'add':
            return inputs + embeddings
        elif self.merge_mode == 'mul':
            return inputs * (embeddings + 1.0)
        elif self.merge_mode == 'zero':
            return embeddings
        else:
            if not self.custom_position_ids:
                embeddings = K.tile(embeddings, [batch_size, 1, 1])
            return K.concatenate([inputs, embeddings])

    def compute_output_shape(self, input_shape):
        if self.custom_position_ids:
            input_shape = input_shape[0]

        if self.merge_mode in ['add', 'mul', 'zero']:
            return input_shape[:2] + (self.output_dim,)
        else:
            return input_shape[:2] + (input_shape[2] + self.output_dim,)

    def get_config(self):
        config = {
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'merge_mode': self.merge_mode,
            'hierarchical': self.hierarchical,
            'embeddings_initializer':
                initializers.serialize(self.embeddings_initializer),
            'custom_position_ids': self.custom_position_ids,
        }
        base_config = super(PositionEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class SinusoidalPositionEmbedding(Layer):
    """定义Sin-Cos位置Embedding
    """

    def __init__(
            self, output_dim, merge_mode='add', custom_position_ids=False, **kwargs
    ):
        super(SinusoidalPositionEmbedding, self).__init__(**kwargs)
        self.output_dim = output_dim
        self.merge_mode = merge_mode
        self.custom_position_ids = custom_position_ids

    def call(self, inputs):
        """如果custom_position_ids，那么第二个输入为自定义的位置id
        """
        if self.custom_position_ids:
            seq_len = K.shape(inputs)[1]
            inputs, position_ids = inputs
            if 'float' not in K.dtype(position_ids):
                position_ids = K.cast(position_ids, K.floatx())
        else:
            input_shape = K.shape(inputs)
            batch_size, seq_len = input_shape[0], input_shape[1]
            position_ids = K.arange(0, seq_len, dtype=K.floatx())[None]

        indices = K.arange(0, self.output_dim // 2, dtype=K.floatx())
        indices = K.pow(10000.0, -2 * indices / self.output_dim)
        embeddings = tf.einsum('bn,d->bnd', position_ids, indices)
        embeddings = K.stack([K.sin(embeddings), K.cos(embeddings)], axis=-1)
        embeddings = K.reshape(embeddings, (-1, seq_len, self.output_dim))

        if self.merge_mode == 'add':
            return inputs + embeddings
        elif self.merge_mode == 'mul':
            return inputs * (embeddings + 1.0)
        elif self.merge_mode == 'zero':
            return embeddings
        else:
            if not self.custom_position_ids:
                embeddings = K.tile(embeddings, [batch_size, 1, 1])
            return K.concatenate([inputs, embeddings])

    def compute_output_shape(self, input_shape):
        if self.custom_position_ids:
            input_shape = input_shape[0]

        if self.merge_mode in ['add', 'mul', 'zero']:
            return input_shape[:2] + (self.output_dim,)
        else:
            return input_shape[:2] + (input_shape[2] + self.output_dim,)

    def get_config(self):
        config = {
            'output_dim': self.output_dim,
            'merge_mode': self.merge_mode,
            'custom_position_ids': self.custom_position_ids,
        }
        base_config = super(SinusoidalPositionEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class RelativePositionEmbedding(Layer):
    """相对位置编码
    来自论文：https://arxiv.org/abs/1803.02155
    """

    def __init__(
            self, input_dim, output_dim, embeddings_initializer='zeros', **kwargs
    ):
        super(RelativePositionEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.embeddings_initializer = initializers.get(embeddings_initializer)

    def build(self, input_shape):
        super(RelativePositionEmbedding, self).build(input_shape)
        self.embeddings = self.add_weight(
            name='embeddings',
            shape=(self.input_dim, self.output_dim),
            initializer=self.embeddings_initializer,
        )

    def call(self, inputs):
        pos_ids = self.compute_position_ids(inputs)
        return K.gather(self.embeddings, pos_ids)

    def compute_position_ids(self, inputs):
        q, v = inputs
        # 计算位置差
        q_idxs = K.arange(0, K.shape(q)[1], dtype='int32')
        q_idxs = K.expand_dims(q_idxs, 1)
        v_idxs = K.arange(0, K.shape(v)[1], dtype='int32')
        v_idxs = K.expand_dims(v_idxs, 0)
        pos_ids = v_idxs - q_idxs
        # 后处理操作
        max_position = (self.input_dim - 1) // 2
        pos_ids = K.clip(pos_ids, -max_position, max_position)
        pos_ids = pos_ids + max_position
        return pos_ids

    def compute_output_shape(self, input_shape):
        return (None, None, self.output_dim)

    def compute_mask(self, inputs, mask):
        return mask[0]

    def get_config(self):
        config = {
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'embeddings_initializer':
                initializers.serialize(self.embeddings_initializer),
        }
        base_config = super(RelativePositionEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class RelativePositionEmbeddingT5(RelativePositionEmbedding):
    """Google T5的相对位置编码
    来自论文：https://arxiv.org/abs/1910.10683
    """

    def __init__(
            self,
            input_dim,
            output_dim,
            max_distance=128,
            bidirectional=True,
            embeddings_initializer='zeros',
            **kwargs
    ):
        super(RelativePositionEmbeddingT5,
              self).__init__(input_dim, output_dim, **kwargs)
        self.max_distance = max_distance
        self.bidirectional = bidirectional

    def compute_position_ids(self, inputs):
        """T5的相对位置分桶（直接翻译自官方T5源码）
        """
        q, v = inputs
        # 计算位置差
        q_idxs = K.arange(0, K.shape(q)[1], dtype='int32')
        q_idxs = K.expand_dims(q_idxs, 1)
        v_idxs = K.arange(0, K.shape(v)[1], dtype='int32')
        v_idxs = K.expand_dims(v_idxs, 0)
        pos_ids = v_idxs - q_idxs
        # 后处理操作
        num_buckets, max_distance = self.input_dim, self.max_distance
        ret = 0
        n = -pos_ids
        if self.bidirectional:
            num_buckets //= 2
            ret += K.cast(K.less(n, 0), 'int32') * num_buckets
            n = K.abs(n)
        else:
            n = K.maximum(n, 0)
        # now n is in the range [0, inf)
        max_exact = num_buckets // 2
        is_small = K.less(n, max_exact)
        val_if_large = max_exact + K.cast(
            K.log(K.cast(n, K.floatx()) / max_exact) /
            np.log(max_distance / max_exact) * (num_buckets - max_exact),
            'int32',
        )
        val_if_large = K.minimum(val_if_large, num_buckets - 1)
        ret += K.switch(is_small, n, val_if_large)
        return ret

    def get_config(self):
        config = {
            'max_distance': self.max_distance,
            'bidirectional': self.bidirectional,
        }
        base_config = super(RelativePositionEmbeddingT5, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class FeedForward(Layer):
    """FeedForward层
    如果activation不是一个list，那么它就是两个Dense层的叠加；如果activation是
    一个list，那么第一个Dense层将会被替换成门控线性单元（Gated Linear Unit）。
    参考论文: https://arxiv.org/abs/2002.05202
    """

    def __init__(
            self,
            units,
            activation='relu',
            use_bias=True,
            kernel_initializer='glorot_uniform',
            **kwargs
    ):
        super(FeedForward, self).__init__(**kwargs)
        self.units = units
        if not isinstance(activation, list):
            activation = [activation]
        self.activation = [activations.get(act) for act in activation]
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)

    @integerize_shape
    def build(self, input_shape):
        super(FeedForward, self).build(input_shape)
        output_dim = input_shape[-1]

        for i, activation in enumerate(self.activation):
            i_dense = Dense(
                units=self.units,
                activation=activation,
                use_bias=self.use_bias,
                kernel_initializer=self.kernel_initializer
            )
            setattr(self, 'i%s_dense' % i, i_dense)

        self.o_dense = Dense(
            units=output_dim,
            use_bias=self.use_bias,
            kernel_initializer=self.kernel_initializer
        )

    @recompute_grad
    def call(self, inputs):
        x = self.i0_dense(inputs)
        for i in range(1, len(self.activation)):
            x = x * getattr(self, 'i%s_dense' % i)(inputs)
        x = self.o_dense(x)
        return x

    def get_config(self):
        config = {
            'units': self.units,
            'activation': [
                activations.serialize(act) for act in self.activation
            ],
            'use_bias': self.use_bias,
            'kernel_initializer':
                initializers.serialize(self.kernel_initializer),
        }
        base_config = super(FeedForward, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class ConditionalRandomField(Layer):
    """纯Keras实现CRF层
    CRF层本质上是一个带训练参数的loss计算层。
    """

    def __init__(self, lr_multiplier=1, **kwargs):
        super(ConditionalRandomField, self).__init__(**kwargs)
        self.lr_multiplier = lr_multiplier  # 当前层学习率的放大倍数

    @integerize_shape
    def build(self, input_shape):
        super(ConditionalRandomField, self).build(input_shape)
        output_dim = input_shape[-1]
        self._trans = self.add_weight(
            name='trans',
            shape=(output_dim, output_dim),
            initializer='glorot_uniform',
            trainable=True
        )
        if self.lr_multiplier != 1:
            K.set_value(self._trans, K.eval(self._trans) / self.lr_multiplier)

    @property
    def trans(self):
        if self.lr_multiplier != 1:
            return self.lr_multiplier * self._trans
        else:
            return self._trans

    def compute_mask(self, inputs, mask=None):
        return None

    def call(self, inputs, mask=None):
        return sequence_masking(inputs, mask, '-inf', 1)

    def target_score(self, y_true, y_pred):
        """计算目标路径的相对概率（还没有归一化）
        要点：逐标签得分，加上转移概率得分。
        """
        point_score = tf.einsum('bni,bni->b', y_true, y_pred)  # 逐标签得分
        trans_score = tf.einsum(
            'bni,ij,bnj->b', y_true[:, :-1], self.trans, y_true[:, 1:]
        )  # 标签转移得分
        return point_score + trans_score

    def log_norm_step(self, inputs, states):
        """递归计算归一化因子
        要点：1、递归计算；2、用logsumexp避免溢出。
        """
        inputs, mask = inputs[:, :-1], inputs[:, -1:]
        states = K.expand_dims(states[0], 2)  # (batch_size, output_dim, 1)
        trans = K.expand_dims(self.trans, 0)  # (1, output_dim, output_dim)
        outputs = tf.reduce_logsumexp(
            states + trans, 1
        )  # (batch_size, output_dim)
        outputs = outputs + inputs
        outputs = mask * outputs + (1 - mask) * states[:, :, 0]
        return outputs, [outputs]

    def dense_loss(self, y_true, y_pred):
        """y_true需要是one hot形式
        """
        # 导出mask并转换数据类型
        mask = K.all(K.greater(y_pred, -1e6), axis=2, keepdims=True)
        mask = K.cast(mask, K.floatx())
        # 计算目标分数
        y_true, y_pred = y_true * mask, y_pred * mask
        target_score = self.target_score(y_true, y_pred)
        # 递归计算log Z
        init_states = [y_pred[:, 0]]
        y_pred = K.concatenate([y_pred, mask], axis=2)
        input_length = K.int_shape(y_pred[:, 1:])[1]
        log_norm, _, _ = K.rnn(
            self.log_norm_step,
            y_pred[:, 1:],
            init_states,
            input_length=input_length
        )  # 最后一步的log Z向量
        log_norm = tf.reduce_logsumexp(log_norm, 1)  # logsumexp得标量
        # 计算损失 -log p
        return log_norm - target_score

    def sparse_loss(self, y_true, y_pred):
        """y_true需要是整数形式（非one hot）
        """
        # y_true需要重新明确一下shape和dtype
        y_true = K.reshape(y_true, K.shape(y_pred)[:-1])
        y_true = K.cast(y_true, 'int32')
        # 转为one hot
        y_true = K.one_hot(y_true, K.shape(self.trans)[0])
        return self.dense_loss(y_true, y_pred)

    def dense_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是one hot形式
        """
        y_true = K.argmax(y_true, 2)
        return self.sparse_accuracy(y_true, y_pred)

    def sparse_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是整数形式（非one hot）
        """
        # 导出mask并转换数据类型
        mask = K.all(K.greater(y_pred, -1e6), axis=2)
        mask = K.cast(mask, K.floatx())
        # y_true需要重新明确一下shape和dtype
        y_true = K.reshape(y_true, K.shape(y_pred)[:-1])
        y_true = K.cast(y_true, 'int32')
        # 逐标签取最大来粗略评测训练效果
        y_pred = K.cast(K.argmax(y_pred, 2), 'int32')
        isequal = K.cast(K.equal(y_true, y_pred), K.floatx())
        return K.sum(isequal * mask) / K.sum(mask)

    def get_config(self):
        config = {
            'lr_multiplier': self.lr_multiplier,
        }
        base_config = super(ConditionalRandomField, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class MaximumEntropyMarkovModel(Layer):
    """（双向）最大熵隐马尔可夫模型
    作用和用法都类似CRF，但是比CRF更快更简单。
    """

    def __init__(self, lr_multiplier=1, hidden_dim=None, **kwargs):
        super(MaximumEntropyMarkovModel, self).__init__(**kwargs)
        self.lr_multiplier = lr_multiplier  # 当前层学习率的放大倍数
        self.hidden_dim = hidden_dim  # 如果非None，则将转移矩阵低秩分解

    @integerize_shape
    def build(self, input_shape):
        super(MaximumEntropyMarkovModel, self).build(input_shape)
        output_dim = input_shape[-1]

        if self.hidden_dim is None:
            self._trans = self.add_weight(
                name='trans',
                shape=(output_dim, output_dim),
                initializer='glorot_uniform',
                trainable=True
            )
            if self.lr_multiplier != 1:
                K.set_value(
                    self._trans,
                    K.eval(self._trans) / self.lr_multiplier
                )
        else:
            self._l_trans = self.add_weight(
                name='l_trans',
                shape=(output_dim, self.hidden_dim),
                initializer='glorot_uniform',
                trainable=True
            )
            self._r_trans = self.add_weight(
                name='r_trans',
                shape=(output_dim, self.hidden_dim),
                initializer='glorot_uniform',
                trainable=True
            )

            if self.lr_multiplier != 1:
                K.set_value(
                    self._l_trans,
                    K.eval(self._l_trans) / self.lr_multiplier
                )
                K.set_value(
                    self._r_trans,
                    K.eval(self._r_trans) / self.lr_multiplier
                )

    @property
    def trans(self):
        if self.lr_multiplier != 1:
            return self.lr_multiplier * self._trans
        else:
            return self._trans

    @property
    def l_trans(self):
        if self.lr_multiplier != 1:
            return self.lr_multiplier * self._l_trans
        else:
            return self._l_trans

    @property
    def r_trans(self):
        if self.lr_multiplier != 1:
            return self.lr_multiplier * self._r_trans
        else:
            return self._r_trans

    def compute_mask(self, inputs, mask=None):
        return None

    def call(self, inputs, mask=None):
        return sequence_masking(inputs, mask, '-inf', 1)

    def reverse_sequence(self, inputs, mask=None):
        if mask is None:
            return [x[:, ::-1] for x in inputs]
        else:
            length = K.cast(K.sum(mask, 1), 'int32')
            return [tf.reverse_sequence(x, length, seq_axis=1) for x in inputs]

    def basic_loss(self, y_true, y_pred, go_backwards=False):
        """y_true需要是整数形式（非one hot）
        """
        # 导出mask并转换数据类型
        mask = K.all(K.greater(y_pred, -1e6), axis=2)
        mask = K.cast(mask, K.floatx())
        # y_true需要重新明确一下shape和dtype
        y_true = K.reshape(y_true, K.shape(y_pred)[:-1])
        y_true = K.cast(y_true, 'int32')
        # 反转相关
        if self.hidden_dim is None:
            if go_backwards:  # 是否反转序列
                y_true, y_pred = self.reverse_sequence([y_true, y_pred], mask)
                trans = K.transpose(self.trans)
            else:
                trans = self.trans
            histoty = K.gather(trans, y_true)
        else:
            if go_backwards:  # 是否反转序列
                y_true, y_pred = self.reverse_sequence([y_true, y_pred], mask)
                r_trans, l_trans = self.l_trans, self.r_trans
            else:
                l_trans, r_trans = self.l_trans, self.r_trans
            histoty = K.gather(l_trans, y_true)
            histoty = tf.einsum('bnd,kd->bnk', histoty, r_trans)
        # 计算loss
        histoty = K.concatenate([y_pred[:, :1], histoty[:, :-1]], 1)
        y_pred = (y_pred + histoty) / 2
        loss = K.sparse_categorical_crossentropy(
            y_true, y_pred, from_logits=True
        )
        return K.sum(loss * mask) / K.sum(mask)

    def sparse_loss(self, y_true, y_pred):
        """y_true需要是整数形式（非one hot）
        """
        loss = self.basic_loss(y_true, y_pred, False)
        loss = loss + self.basic_loss(y_true, y_pred, True)
        return loss / 2

    def dense_loss(self, y_true, y_pred):
        """y_true需要是one hot形式
        """
        y_true = K.argmax(y_true, 2)
        return self.sparse_loss(y_true, y_pred)

    def basic_accuracy(self, y_true, y_pred, go_backwards=False):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是整数形式（非one hot）
        """
        # 导出mask并转换数据类型
        mask = K.all(K.greater(y_pred, -1e6), axis=2)
        mask = K.cast(mask, K.floatx())
        # y_true需要重新明确一下shape和dtype
        y_true = K.reshape(y_true, K.shape(y_pred)[:-1])
        y_true = K.cast(y_true, 'int32')
        # 反转相关
        if self.hidden_dim is None:
            if go_backwards:  # 是否反转序列
                y_true, y_pred = self.reverse_sequence([y_true, y_pred], mask)
                trans = K.transpose(self.trans)
            else:
                trans = self.trans
            histoty = K.gather(trans, y_true)
        else:
            if go_backwards:  # 是否反转序列
                y_true, y_pred = self.reverse_sequence([y_true, y_pred], mask)
                r_trans, l_trans = self.l_trans, self.r_trans
            else:
                l_trans, r_trans = self.l_trans, self.r_trans
            histoty = K.gather(l_trans, y_true)
            histoty = tf.einsum('bnd,kd->bnk', histoty, r_trans)
        # 计算逐标签accuracy
        histoty = K.concatenate([y_pred[:, :1], histoty[:, :-1]], 1)
        y_pred = (y_pred + histoty) / 2
        y_pred = K.cast(K.argmax(y_pred, 2), 'int32')
        isequal = K.cast(K.equal(y_true, y_pred), K.floatx())
        return K.sum(isequal * mask) / K.sum(mask)

    def sparse_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是整数形式（非one hot）
        """
        accuracy = self.basic_accuracy(y_true, y_pred, False)
        accuracy = accuracy + self.basic_accuracy(y_true, y_pred, True)
        return accuracy / 2

    def dense_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是one hot形式
        """
        y_true = K.argmax(y_true, 2)
        return self.sparse_accuracy(y_true, y_pred)

    def get_config(self):
        config = {
            'lr_multiplier': self.lr_multiplier,
            'hidden_dim': self.hidden_dim,
        }
        base_config = super(MaximumEntropyMarkovModel, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class GlobalPointer(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, RoPE=True, **kwargs):
        super(GlobalPointer, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE

    def build(self, input_shape):
        super(GlobalPointer, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, mask=None):
        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class GlobalPointerDropout(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, RoPE=True, **kwargs):
        super(GlobalPointerDropout, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE

    def build(self, input_shape):
        super(GlobalPointerDropout, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        self.dp = Dropout(0.1)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, mask=None):
        # 输入变换
        inputs = self.dp(inputs)
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerDropout, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class CNNEncoder(Layer):
    """全局指针模块,加上多特征融合
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self,
                 embedding_dim: int,
                 num_filters: int,
                 ngram_filter_sizes,
                 conv_layer_activation=None,
                 output_dim=None,
                 **kwargs
                 ):
        super(CNNEncoder, self).__init__(**kwargs)
        self._embedding_dim = embedding_dim
        self._num_filters = num_filters
        self._ngram_filter_sizes = ngram_filter_sizes
        self._activation = conv_layer_activation or Activation('relu')
        self._output_dim = output_dim



    def build(self, input_shape):
        print('-----{}'.format(input_shape))
        super(CNNEncoder, self).build(input_shape)
        self._convolution_layers = [
            Conv2D(
                filters=self._num_filters,
                kernel_size=(1, ngram_size),
            )
            for ngram_size in self._ngram_filter_sizes
        ]
        # for i, conv_layer in enumerate(self._convolution_layers):
        #     self.add_module("conv_layer_%d" % i, conv_layer)

        maxpool_output_dim = self._num_filters * len(self._ngram_filter_sizes)
        if self._output_dim:
            self.projection_layer = Dense(self._output_dim)
        else:
            self.projection_layer = None
            self._output_dim = maxpool_output_dim

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, tokens, mask=None):
        print(K.int_shape(tokens))

        # Our input is expected to have shape `(batch_size, num_tokens, embedding_dim)`.  The
        # convolution layers expect input of shape `(batch_size, in_channels, sequence_length)`,
        # where the conv layer `in_channels` is our `embedding_dim`.  We thus need to transpose the
        # tensor first.

        # Each convolution layer returns output of size `(batch_size, num_filters, pool_length)`,
        # where `pool_length = num_tokens - ngram_size + 1`.  We then do an activation function,
        # masking, then do max pooling over each filter for the whole input sequence.
        # Because our max pooling is simple, we just use `torch.max`.  The resultant tensor has shape
        # `(batch_size, num_conv_layers * num_filters)`, which then gets projected using the
        # projection layer, if requested.

        # To ensure the cnn_encoder respects masking we add a large negative value to
        # the activations of all filters that convolved over a masked token. We do this by
        # first enumerating all filters for a given convolution size (torch.arange())
        # then by comparing it to an index of the last filter that does not involve a masked
        # token (.ge()) and finally adjusting dimensions to allow for addition and multiplying
        # by a large negative value (.unsqueeze())
        filter_outputs = []

        # shape: (batch_size, 1)
        for i in range(len(self._convolution_layers)):
            convolution_layer = self._convolution_layers[i]
            # Forward pass of the convolutions.
            # shape: (batch_size, num_filters, pool_length)
            activations = self._activation(convolution_layer(tokens))

            # Create activation mask.
            # shape: (batch_size, pool_length)

            # shape: (batch_size, pool_length)

            # shape: (batch_size, num_filters, pool_length)


            # Replace masked out values with smallest possible value of the dtype so
            # that max pooling will ignore these activations.
            # shape: (batch_size, pool_length)

            # activations = activations + (activations_mask * min_value_of_dtype(activations.dtype))
            activationsmax = tf.reduce_max(activations, axis=2, keepdims=False)
            # Pick out the max filters
            filter_outputs.append(activationsmax)

        # Now we have a list of `num_conv_layers` tensors of shape `(batch_size, num_filters)`.
        # Concatenating them gives us a tensor of shape `(batch_size, num_filters * num_conv_layers)`.
        maxpool_output = (
            K.concatenate(filter_outputs, axis=-1) if len(filter_outputs) > 1 else filter_outputs[0]
        )

        # Replace the maxpool activations that picked up the masks with 0s
        # maxpool_output[maxpool_output == min_value_of_dtype(maxpool_output.dtype)] = 0.0

        if self.projection_layer:
            result = self.projection_layer(maxpool_output)
        else:
            result = maxpool_output
        return result

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1], self._output_dim)

    def get_config(self):
        config = {
            'outputdim': self._output_dim,
        }
        base_config = super(CNNEncoder, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

def backend_zixing_reshape(x):
    K.reshape(x, (-1, 80, 384))

def backend_pinyin_reshape(x):
    K.reshape(x, (-1, 32, 384))


class GlobalPointerMF(Layer):
    """全局指针模块,加上多特征融合
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, pinyin_size, zixing_size, RoPE=True, **kwargs):
        super(GlobalPointerMF, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE
        self.pinyin_size = pinyin_size
        self.zixing_size = zixing_size

    def build(self, input_shape):
        super(GlobalPointerMF, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)

        self.pinyinEmbedding = Embedding(self.pinyin_size, 384)
        # self.zixingEmbedding = Embedding(self.zixing_size, 384)
        self.zixingEmbedding = Embedding(self.zixing_size, 384)
        self.pinyinRP = Lambda(backend_pinyin_reshape)
        self.zixingRP = Lambda(backend_zixing_reshape)

        # self.pinyinEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(2, 3, 4, 5), output_dim=256)
        # self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=256)
        self.pinyinEncoder = Conv1D(384, kernel_size=3)
        self.pinyinPooler = MaxPool1D(256)
        self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=256)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, pinyin, zixing, mask=None):
        # 特征拼接
        print(K.int_shape(pinyin))
        pinyinemb = self.pinyinEmbedding(pinyin) # batch, seqlen, pylen, embsize
        zixingemb = self.zixingEmbedding(zixing) # batch, seqlen, zxlen, embsize

        # seqlen, pylen, embsize1 = K.shape(pinyinemb)
        # seqlen, zxlen, embsize2 = K.shape(zixingemb)
        print(K.int_shape(pinyinemb))
        pinyinemb = tf.reduce_sum(pinyinemb, axis=2, keepdims=False)
        zixingemb = tf.reduce_sum(zixingemb, axis=2, keepdims=False)
        print(K.int_shape(pinyinemb))

        # pinyinemb_rp = self.pinyinRP(pinyinemb)
        # zixingemb_rp = self.zixingRP(zixingemb)
        #
        # encoded_pinyin = self.pinyinEncoder(pinyinemb_rp) # batch * seqlen, outdim
        # encoded_zixing = self.zixingEncoder(zixingemb_rp) # batch * seqlen, outdim
        #
        # print(encoded_pinyin)

        # encoded_pinyin_rp = K.reshape(encoded_pinyin, (seqlen, -1))
        # encoded_zixing_rp = K.reshape(encoded_zixing, (seqlen, -1))

        inputs = K.concatenate([inputs, pinyinemb, zixingemb], axis=-1)
        # inputs = K.concatenate([inputs, pinyinemb, zixingemb], axis=-1)

        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerMF, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class GlobalPointerMFV2(Layer):
    """全局指针模块,加上多特征融合,针对不同特征，利用卷积+池化来提取特征
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, pinyin_size, zixing_size, RoPE=True, **kwargs):
        super(GlobalPointerMFV2, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE
        self.pinyin_size = pinyin_size
        self.zixing_size = zixing_size

    def build(self, input_shape):
        super(GlobalPointerMFV2, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        from keras.layers import Embedding
        self.pinyinEmbedding = Embedding(input_dim=self.pinyin_size, output_dim=384)
        self.zixingEmbedding = Embedding(self.zixing_size, 384)
        # self.zixingEmbedding = Embedding(1536, 384)
        self.pinyinRP = Lambda(backend_pinyin_reshape)
        self.zixingRP = Lambda(backend_zixing_reshape)

        self.pinyinEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(2, 3, 4, 5), output_dim=384)
        self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=384)
        # self.pinyinEncoder = Conv2D(filters=144, kernel_size=(1, 3)) # 输入是batch， seqlen， pylen， embsize
        # self.pinyinPooler = MaxPool1D(256)
        # self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=256)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, pinyin, zixing, mask=None):
        # 特征拼接
        print(K.int_shape(pinyin))
        pinyinemb = self.pinyinEmbedding(pinyin) # batch, seqlen, pylen, embsize
        zixingemb = self.zixingEmbedding(zixing) # batch, seqlen, zxlen, embsize
        print(pinyinemb)
        print(pinyinemb.get_shape())
        print(pinyinemb._keras_shape)
        # seqlen, pylen, embsize1 = K.shape(pinyinemb)
        # seqlen, zxlen, embsize2 = K.shape(zixingemb)

        encoded_pinyin = self.pinyinEncoder(pinyinemb) # batch * seqlen, outdim
        encoded_zixing = self.zixingEncoder(zixingemb) # batch * seqlen, outdim


        # encoded_pinyin_rp = K.reshape(encoded_pinyin, (seqlen, -1))
        # encoded_zixing_rp = K.reshape(encoded_zixing, (seqlen, -1))

        inputs = K.concatenate([inputs, encoded_pinyin, encoded_zixing], axis=-1)
        # inputs = K.concatenate([inputs, pinyinemb, zixingemb], axis=-1)

        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerMFV2, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class GlobalPointerMFback(Layer):
    """全局指针模块,加上多特征融合
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, pinyin_size, zixing_size, RoPE=True, **kwargs):
        super(GlobalPointerMF, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE
        self.pinyin_size = pinyin_size
        self.zixing_size = zixing_size

    def build(self, input_shape):
        super(GlobalPointerMF, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        self.pinyinEmbedding = Embedding(self.pinyin_size, 384)
        self.zixingEmbedding = Embedding(self.zixing_size, 384)

        self.pinyinRP = Lambda(backend_pinyin_reshape)
        self.zixingRP = Lambda(backend_zixing_reshape)

        # self.pinyinEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(2, 3, 4, 5), output_dim=256)
        # self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=256)
        self.pinyinEncoder = Conv1D(384, kernel_size=3)
        self.pinyinPooler = MaxPool1D(256)
        self.zixingEncoder = CNNEncoder(embedding_dim=256, num_filters=128, ngram_filter_sizes=(3, 4, 5, 6), output_dim=256)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, pinyin, zixing, mask=None):
        # 特征拼接
        print(K.int_shape(pinyin))
        pinyinemb = self.pinyinEmbedding(pinyin) # batch, seqlen, pylen, embsize
        zixingemb = self.zixingEmbedding(zixing) # batch, seqlen, zxlen, embsize

        # seqlen, pylen, embsize1 = K.shape(pinyinemb)
        # seqlen, zxlen, embsize2 = K.shape(zixingemb)
        print(K.int_shape(pinyinemb))
        pinyinemb_rp = self.pinyinRP(pinyinemb)
        zixingemb_rp = self.zixingRP(zixingemb)

        encoded_pinyin = self.pinyinEncoder(pinyinemb_rp) # batch * seqlen, outdim
        encoded_zixing = self.zixingEncoder(zixingemb_rp) # batch * seqlen, outdim

        print(encoded_pinyin)

        # encoded_pinyin_rp = K.reshape(encoded_pinyin, (seqlen, -1))
        # encoded_zixing_rp = K.reshape(encoded_zixing, (seqlen, -1))

        inputs = K.concatenate([inputs, encoded_pinyin, encoded_zixing], axis=-1)
        # inputs = K.concatenate([inputs, pinyinemb, zixingemb], axis=-1)

        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerMF, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


class GlobalPointerPseudo(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    伪标签动态调整权重
    """

    def __init__(self, heads, head_size, RoPE=True, **kwargs):
        super(GlobalPointerPseudo, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE

    def build(self, input_shape):
        super(GlobalPointerPseudo, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        self.lossweight = self.add_weight(name='lossweight',
                                      shape=(1,),  # 假设输入tensor只有一个维度（不算batch的维度）
                                    initializer=initializers.get('uniform'),
                                      trainable=True)  # 如果要定义可训练参数这里一定要选择True

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, ispseudo=[], mask=None):
        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12

        weight = K.sigmoid(self.lossweight)
        # bpseudo = K.reshape(ispseudo, [-1])
        bpseudo = K.reshape(ispseudo, [-1, 1])
        # scale返回
        return [logits / self.head_size ** 0.5, bpseudo]

    def compute_output_shape(self, input_shape):
        return [(input_shape[0], self.heads, input_shape[1], input_shape[1]), (input_shape[0], 1)]

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerPseudo, self).get_config()
        return dict(list(GlobalPointerPseudo.items()) + list(config.items()))

class GlobalPointerTagger(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, RoPE=True, **kwargs):
        super(GlobalPointerTagger, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE

    def build(self, input_shape):
        super(GlobalPointerTagger, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, mask=None):
        # 输入变换
        inputs = self.dense(inputs) # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1) # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2) # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:] # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs) # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw) # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size**0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerTagger, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class GlobalPointerKG(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, maxlen, cidian_dim, is_kg, RoPE=True, **kwargs):
        super(GlobalPointerKG, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE
        self.maxlen = maxlen
        self.cidian_dim = cidian_dim
        self.is_kg = is_kg

    def build(self, input_shape):
        super(GlobalPointerKG, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        self._cidian_word_embeddings = Embedding(self.maxlen, self.cidian_dim)

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, words, mask=None):
        # inputs=>b, seqlen, hidden
        if self.is_kg:
            for i, word in enumerate(words): # b numadd l
                # l, l
                cidian_word_embedding = self._cidian_word_embeddings(word)
                if i == 0:
                    cidian_word_embeddings = cidian_word_embedding
                else:
                    cidian_word_embeddings = tf.concat([cidian_word_embeddings, cidian_word_embedding], axis=-1)
            inputs = tf.concat([inputs, cidian_word_embeddings], axis=-1)

        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerKG, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))



class GlobalPointerKGALL(Layer):
    """全局指针模块
    将序列的每个(start, end)作为整体来进行判断
    """

    def __init__(self, heads, head_size, maxlen, cidian_dim, is_kg, RoPE=True, **kwargs):
        super(GlobalPointerKGALL, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.RoPE = RoPE
        self.maxlen = maxlen
        self.cidian_dim = cidian_dim
        self.is_kg = is_kg

    def build(self, input_shape):
        super(GlobalPointerKGALL, self).build(input_shape)
        self.dense = Dense(self.head_size * self.heads * 2)
        self._cidian_word_embeddings_lst = []
        for _ in range(self.heads): # 类别数量
            self._cidian_word_embeddings_lst.append(Embedding(2, self.cidian_dim))

    def compute_mask(self, inputs, mask=None):
        return None

    @recompute_grad
    def call(self, inputs, words, mask=None):
        if self.is_kg:
            for i, kgword in enumerate(words): # 10维 因为外围传递10个输入进来
                # 内部维度为batch ， kglabels， seqlen
                kglst = K.permute_dimensions(kgword, (1, 0, 2))
                # kglst = tf.split(kgword, self.heads, axis=1)
                # print(kglst[0].shape)
                for j in range(self.heads): # heads
                    kgclsword = kglst[j]
                    # kgclsword maxlen
                    cidian_word_embedding_layer = self._cidian_word_embeddings_lst[j]
                    if j == 0:
                        cidian_word_embedding = cidian_word_embedding_layer(kgclsword)
                    else:
                        cidian_word_embedding += cidian_word_embedding_layer(kgclsword)

                if i == 0:
                    cidian_word_embeddings = cidian_word_embedding
                else:
                    cidian_word_embeddings = tf.concat([cidian_word_embeddings, cidian_word_embedding], axis=-1)
            inputs = tf.concat([inputs, cidian_word_embeddings], axis=-1)

        # 输入变换
        inputs = self.dense(inputs)  # b, l,  self.head_size * self.heads * 2
        inputs = tf.split(inputs, self.heads, axis=-1)  # 划分多头 self.heads, b, l, self.head_size * 2
        inputs = K.stack(inputs, axis=-2)  # 将heads 叠加一起 b, l, self.heads, self.head_size * 2
        qw, kw = inputs[..., :self.head_size], inputs[..., self.head_size:]  # 一个start，一个end
        # qw b,l,heads,headsize   kw b,l,heads, headsize
        # RoPE编码
        if self.RoPE:
            pos = SinusoidalPositionEmbedding(self.head_size, 'zero')(inputs)  # b, l, headsize
            cos_pos = K.repeat_elements(pos[..., None, 1::2], 2, -1)
            sin_pos = K.repeat_elements(pos[..., None, ::2], 2, -1)
            qw2 = K.stack([-qw[..., 1::2], qw[..., ::2]], 4)
            qw2 = K.reshape(qw2, K.shape(qw))
            qw = qw * cos_pos + qw2 * sin_pos
            kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
            kw2 = K.reshape(kw2, K.shape(kw))
            kw = kw * cos_pos + kw2 * sin_pos
        # 计算内积
        logits = tf.einsum('bmhd,bnhd->bhmn', qw, kw)  # b, heads, l, l
        # 排除padding
        logits = sequence_masking(logits, mask, '-inf', 2)
        logits = sequence_masking(logits, mask, '-inf', 3)
        # 排除下三角
        mask = tf.matrix_band_part(K.ones_like(logits), 0, -1)
        logits = logits - (1 - mask) * 1e12
        # scale返回
        return logits / self.head_size ** 0.5

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.heads, input_shape[1], input_shape[1])

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'RoPE': self.RoPE,
        }
        base_config = super(GlobalPointerKGALL, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class Loss(Layer):
    """特殊的层，用来定义复杂loss
    """

    def __init__(self, output_axis=None, **kwargs):
        super(Loss, self).__init__(**kwargs)
        self.output_axis = output_axis

    def call(self, inputs, mask=None):
        loss = self.compute_loss(inputs, mask)
        self.add_loss(loss, inputs=inputs)
        if self.output_axis is None:
            return inputs
        elif isinstance(self.output_axis, list):
            return [inputs[i] for i in self.output_axis]
        else:
            return inputs[self.output_axis]

    def compute_loss(self, inputs, mask=None):
        raise NotImplementedError

    def compute_output_shape(self, input_shape):
        if self.output_axis is None:
            return input_shape
        elif isinstance(self.output_axis, list):
            return [input_shape[i] for i in self.output_axis]
        else:
            return input_shape[self.output_axis]

    def compute_mask(self, inputs, mask):
        if mask is not None:
            if self.output_axis is None:
                return mask
            elif isinstance(self.output_axis, list):
                return [mask[i] for i in self.output_axis]
            else:
                return mask[self.output_axis]

    def get_config(self):
        config = {
            'output_axis': self.output_axis,
        }
        base_config = super(Loss, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))


custom_objects = {
    'Embedding': Embedding,
    'BiasAdd': BiasAdd,
    'Concatenate1D': Concatenate1D,
    'MultiHeadAttention': MultiHeadAttention,
    'LayerNormalization': LayerNormalization,
    'PositionEmbedding': PositionEmbedding,
    'SinusoidalPositionEmbedding': SinusoidalPositionEmbedding,
    'RelativePositionEmbedding': RelativePositionEmbedding,
    'RelativePositionEmbeddingT5': RelativePositionEmbeddingT5,
    'FeedForward': FeedForward,
    'ConditionalRandomField': ConditionalRandomField,
    'MaximumEntropyMarkovModel': MaximumEntropyMarkovModel,
    'GlobalPointer': GlobalPointer,
    'Loss': Loss,
}

keras.utils.get_custom_objects().update(custom_objects)