Add files via upload

Author: pass-lin

bert4keras3/Layers_add/Attentions.py

@@ -173,7 +173,7 @@ def call(self, inputs, mask=None, **kwargs):
             o = self.o_dense(ops.reshape(o, [b,s,-1]))
         # 返回结果
 
-        
+
         if use_cache:
             return o,cache
         if self.return_attention_scores:
@@ -282,17 +282,18 @@ def pay_attention_to(self, inputs, mask=None, **kwargs):
             a = a * ops.cast(1/np.sqrt(self.key_size), dtype=qw.dtype)
         if a_bias is not None and ops.ndim(a_bias) == 3:
             a_bias = align(a_bias, [0, -2, -1], ops.ndim(a))
-        
-        A,mask = attention_normalize(a, v_mask, -1, self.normalization, a_bias)
+        A = attention_normalize(a, v_mask, -1, self.normalization, a_bias)
 
         if self.attention_dropout:
-            A,mask = self.dropout(A)
+            A = self.dropout(A)
+
         # 完成输出
         if self.query_head!=self.heads:
             o = ops.einsum("bkgts,bskh->btkgh", A, vw)
             o = ops.reshape(o, (b, s, self.query_head, -1))
         else:
             o = ops.einsum('bhjk,bkhd->bjhd', A, vw)
+        
         if p_bias == 'typical_relative':
             o = o + ops.einsum('bhjk,jkd->bjhd', A, position_bias)
 

bert4keras3/Layers_add/FFN.py

@@ -13,6 +13,7 @@ def __init__(
         activation='relu',
         use_bias=True,
         kernel_initializer='glorot_uniform',
+
         **kwargs
     ):
         super(FeedForward, self).__init__(**kwargs)
@@ -98,34 +99,35 @@ class LLamaFeedForward(FeedForward):
     def build(self, input_shape):
         super(FeedForward, self).build(input_shape)
         output_dim = input_shape[-1]
-        self._feedforward_intermediate_dense = keras.layers.Dense(
+        self._feedforward_gate_dense = keras.layers.Dense(
             self.units,
             kernel_initializer=self.kernel_initializer,
             use_bias=self.use_bias,
-            name="feedforward_intermediate_dense",
+            name="feedforward_gate_dense",
         )
-        self._feedforward_gate_dense = keras.layers.Dense(
+        self._feedforward_intermediate_dense = keras.layers.Dense(
             self.units,
             kernel_initializer=self.kernel_initializer,
             use_bias=self.use_bias,
-            name="feedforward_gate_dense",
+            name="feedforward_intermediate_dense",
         )
+        
 
         self._feedforward_output_dense = keras.layers.Dense(
             output_dim,
             kernel_initializer=self.kernel_initializer,
-            use_bias=False,
-            dtype=self.use_bias,
+            use_bias=self.use_bias,
             name="feedforward_output_dense",
         )
     @recompute_grad
     def call(self, x):
+
         activation = activations.get(self.activation[0])
         gate_output = self._feedforward_gate_dense(x)
-        gate_output = ops.cast(gate_output, "float32")
+        #gate_output = ops.cast(gate_output, "float32")
         gate_output = activation(gate_output)
-        gate_output = ops.cast(gate_output, x.dtype)
+        #gate_output = ops.cast(gate_output, x.dtype)
         x = self._feedforward_intermediate_dense(x)
         x = self._feedforward_output_dense(ops.multiply(x, gate_output))
-        return x
+        return x#
 

bert4keras3/Layers_add/LayerNorms.py

@@ -204,6 +204,7 @@ def call(self, x):
         x = ops.cast(x, "float32")
         var = ops.mean(ops.power(x, 2), axis=-1, keepdims=True)
         x = x * ops.rsqrt(var + self.epsilon)
+        
         return ops.cast(x, self.compute_dtype) * self.scale
 
     def get_config(self):

bert4keras3/Layers_add/__pycache__/Attentions.cpython-310.pyc

@@ -7,15 +7,23 @@ def __init__(self, with_lm=True,
                  use_EinsumDense = True,
                  flatten_o_dense=False,
                  use_bias = False,
+                 input_scale =True,
+                 share_emebding=True,
+                 rope_mode='keras',
                  **kwargs):
         super(Gemma, self).__init__(**kwargs)
         self.with_lm = with_lm
         self.max_wavelength = max_wavelength
         self.scaling_factor = scaling_factor
+        self.rope_mode = rope_mode
+        self.share_emebding = share_emebding
         self.use_dense_bias = use_dense_bias
+        self.input_scale = input_scale
         self.flatten_o_dense = flatten_o_dense
         self.use_EinsumDense = use_EinsumDense
         self.use_bias = use_bias
+        self.layer_norm_type = RMSNormalization
+        self.ffn_type = GemmaFeedForward
     def apply_embeddings(self, inputs):
         inputs = inputs[:]
 
@@ -55,12 +63,13 @@ def apply_embeddings(self, inputs):
 
         def mul(x):
             return x * ops.cast(ops.sqrt(self.hidden_size), x.dtype)
-        x = self.apply(
-            inputs=x,
-            layer=Lambda,
-            function=mul,
-            name='Multiply'
-        )
+        if self.input_scale:
+            x = self.apply(
+                inputs=x,
+                layer=Lambda,
+                function=mul,
+                name='Multiply'
+            )
 
         x = self.apply(
             inputs=x,
@@ -92,7 +101,7 @@ def apply_main_layers(self, inputs, index):
 
         x = self.apply(
             inputs=x,
-            layer=RMSNormalization,
+            layer=self.layer_norm_type,
             epsilon=1e-6,
             name='%s-Norm' % attention_name
         )
@@ -137,19 +146,11 @@ def apply_main_layers(self, inputs, index):
 
         x = self.apply(
             inputs=x,
-            layer=RMSNormalization,
+            layer=self.layer_norm_type,
             epsilon=1e-6,
             name='%s-Norm' % feed_forward_name
         )
-        x = self.apply(
-            inputs=x,
-            layer=GemmaFeedForward,
-            units=self.intermediate_size,
-            activation=self.hidden_act,
-            use_bias=self.use_dense_bias,
-            kernel_initializer=self.initializer,
-            name=feed_forward_name
-        )
+        x = self.apply_ffn(x,feed_forward_name)
         x = self.apply(
             inputs=x,
             layer=Dropout,
@@ -170,7 +171,7 @@ def apply_final_layers(self, inputs):
 
         x = self.apply(
             inputs=x,
-            layer=RMSNormalization,
+            layer=self.layer_norm_type,
             epsilon=1e-6,
             name='Output-Norm'
         )
@@ -183,19 +184,42 @@ def apply_final_layers(self, inputs):
 
         if self.with_lm:
             lm_activation = 'softmax' if self.with_lm is True else self.with_lm
-            x = self.apply(
+            if self.share_emebding:
+                x = self.apply(
+                        inputs=x,
+                        layer=Embedding,
+                        arguments={'mode': 'dense'},
+                        name='Embedding-Token'
+                    )
+                x = self.apply(
+                        inputs=x,
+                        layer=Activation,
+                        activation=lm_activation,
+                        name='Output-LM-Activation'
+                    )
+            else:
+                x = self.apply(
                     inputs=x,
-                    layer=Embedding,
-                    arguments={'mode': 'dense'},
-                    name='Embedding-Token'
-                )
-            x = self.apply(
-                    inputs=x,
-                    layer=Activation,
+                    layer=Dense,
+                    units=self.vocab_size,
                     activation=lm_activation,
-                    name='Output-LM-Activation'
+                    use_bias=False,
+                    kernel_initializer=self.initializer,
+                    name='Decoder-Output-LM'
                 )
+
 
+        return x
+    def apply_ffn(self,x,feed_forward_name):
+        x = self.apply(
+            inputs=x,
+            layer=self.ffn_type,
+            units=self.intermediate_size,
+            activation=self.hidden_act,
+            use_bias=self.use_dense_bias,
+            kernel_initializer=self.initializer,
+            name=feed_forward_name
+        )
         return x
     def apply_main_cache_layers(self, inputs, index,self_cache_update_index,
                                 cross_cache_update_index=None,
@@ -243,3 +267,11 @@ def apply_main_cache_layers(self, inputs, index,self_cache_update_index,
 
         return [x,caches]
 
+class Llama(Gemma):
+    def __init__(self, input_scale =False,use_EinsumDense=False,
+                 share_emebding=False,**kwargs):
+        super(Llama, self).__init__(input_scale=input_scale,
+                                    use_EinsumDense=use_EinsumDense,
+                                    share_emebding=share_emebding,**kwargs)
+        self.layer_norm_type = LlamaLayerNorm
+        self.ffn_type = LLamaFeedForward

bert4keras3/Layers_add/__pycache__/Attentions.cpython-311.pyc

@@ -1,6 +1,6 @@
 #! -*- coding: utf-8 -*-
 
-__version__ = '1.1.2'
+__version__ = '1.3'
 
 from bert4keras3 import backend,layers,models,snippets,tokenizers
 from bert4keras3.backend import ops

bert4keras3/Layers_add/__pycache__/Embeddings.cpython-310.pyc

@@ -365,23 +365,23 @@ def attention_normalize(a, mask=None, axis=-1, method='softmax', bias=None):
         att_mask = mask
         for i in range(ops.ndim(a)-ops.ndim(mask)):
             att_mask = ops.expand_dims(att_mask,0)
-        return ops.cast(keras.layers.Softmax(dtype="float32",axis=axis)(a,mask=att_mask),ori_dtype),mask
+        return ops.cast(keras.layers.Softmax(dtype="float32",axis=axis)(a,mask=att_mask),ori_dtype)
     a, mask = sequence_masking(a, mask, -np.inf, axis, bias, True)
 
     if method == 'softmax' :
-        return ops.softmax(a,axis=axis),mask
+        return ops.softmax(a,axis=axis)
     else:
         if mask is None:
             l = ops.cast(ops.shape(a)[-1], keras.mixed_precision.dtype_policy().name)
         else:
             mask = ops.cast(mask, keras.mixed_precision.dtype_policy().name)
             l = ops.sum(mask, axis=axis, keepdims=True)
         if method == 'squared_relu':
-            return ops.relu(a)**2 / l,mask
+            return ops.relu(a)**2 / l
         elif method == 'softmax_plus':
             l = ops.maximum(l, 16)  # 极短序列scale反而不好
-            return ops.softmax(a * ops.log(l) / np.log(512), axis=axis),mask
-    return a,mask
+            return ops.softmax(a * ops.log(l) / np.log(512), axis=axis)
+    return a
 
 
 def sinusoidal_embeddings(pos, dim, base=10000):

bert4keras3/Layers_add/__pycache__/Embeddings.cpython-311.pyc

@@ -93,7 +93,9 @@ def build_transformer_model(
         'mt5.1.1_encoder': T5_Encoder,
         'mt5.1.1_decoder': T5_Decoder,
         'gemma':Gemma,
-
+        'llama':Llama, 
+        'qwen':Llama, 
+        'yi':Llama, 
         'misakat5':MisakaT5,
     }