允许roformer自定义位置id

修复roformer自定义位置id不起作用
chess122 · Mar 30, 2022 · 9769d5e · 9769d5e
2 parents 56407a5 + 1fae360
commit 9769d5e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/bert4keras/layers.py b/bert4keras/layers.py
@@ -837,8 +837,8 @@ def call(self, inputs):
         """如果custom_position_ids，那么第二个输入为自定义的位置id
         """
         if self.custom_position_ids:
-            seq_len = K.shape(inputs)[1]
             inputs, position_ids = inputs
+            seq_len = K.shape(inputs)[1]
             if 'float' not in K.dtype(position_ids):
                 position_ids = K.cast(position_ids, K.floatx())
         else:

diff --git a/bert4keras/models.py b/bert4keras/models.py
@@ -1197,13 +1197,16 @@ def compute_position_bias(self, inputs=None):
         """Sinusoidal位置编码（直接返回）
         """
         if self.position_bias is None:
-
-            x = inputs
+            if self.custom_position_ids:
+                x = [inputs, self.inputs[2]]
+            else:
+                x = inputs
             self.position_bias = self.apply(
                 inputs=x,
                 layer=SinusoidalPositionEmbedding,
                 output_dim=self.attention_key_size,
                 merge_mode='zero',
+                custom_position_ids=self.custom_position_ids,
                 name='Embedding-Rotary-Position'
             )