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):