[Neural Search] change output_emb_size to self.ptm.config.hidden_size (…

…#5127) * change output_emb_size to self.ptm.config['hidden_size'] * Update applications/neural_search/recall/milvus/base_model.py Co-authored-by: Sijun He <sijun.he@hotmail.com> * Update applications/neural_search/recall/milvus/base_model.py Co-authored-by: Sijun He <sijun.he@hotmail.com> * Update applications/neural_search/recall/simcse/model.py Co-authored-by: Sijun He <sijun.he@hotmail.com> * change to self.ptm.config.hidden_size --------- Co-authored-by: Sijun He <sijun.he@hotmail.com>
PaddlePaddle · Mar 7, 2023 · e40e40b · e40e40b
1 parent ee390a6
commit e40e40b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git a/applications/neural_search/recall/in_batch_negative/base_model.py b/applications/neural_search/recall/in_batch_negative/base_model.py
@@ -32,7 +32,9 @@ def __init__(self, pretrained_model, dropout=None, output_emb_size=None):
         self.output_emb_size = output_emb_size
         if output_emb_size > 0:
             weight_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.TruncatedNormal(std=0.02))
-            self.emb_reduce_linear = paddle.nn.Linear(768, output_emb_size, weight_attr=weight_attr)
+            self.emb_reduce_linear = paddle.nn.Linear(
+                self.ptm.config.hidden_size, output_emb_size, weight_attr=weight_attr
+            )
 
     def get_pooled_embedding(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None):
         _, cls_embedding = self.ptm(input_ids, token_type_ids, position_ids, attention_mask)
@@ -95,7 +97,9 @@ def __init__(self, pretrained_model, dropout=None, output_emb_size=None):
         self.output_emb_size = output_emb_size
         if output_emb_size > 0:
             weight_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.TruncatedNormal(std=0.02))
-            self.emb_reduce_linear = paddle.nn.Linear(768, output_emb_size, weight_attr=weight_attr)
+            self.emb_reduce_linear = paddle.nn.Linear(
+                self.ptm.config.hidden_size, output_emb_size, weight_attr=weight_attr
+            )
 
     @paddle.jit.to_static(
         input_spec=[

diff --git a/applications/neural_search/recall/milvus/base_model.py b/applications/neural_search/recall/milvus/base_model.py
@@ -13,9 +13,6 @@
 # limitations under the License.
 
 import abc
-import sys
-
-import numpy as np
 
 import paddle
 import paddle.nn as nn
@@ -35,7 +32,9 @@ def __init__(self, pretrained_model, dropout=None, output_emb_size=None):
         self.output_emb_size = output_emb_size
         if output_emb_size > 0:
             weight_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.TruncatedNormal(std=0.02))
-            self.emb_reduce_linear = paddle.nn.Linear(768, output_emb_size, weight_attr=weight_attr)
+            self.emb_reduce_linear = paddle.nn.Linear(
+                self.ptm.config.hidden_size, output_emb_size, weight_attr=weight_attr
+            )
 
     @paddle.jit.to_static(
         input_spec=[
@@ -106,7 +105,9 @@ def __init__(self, pretrained_model, dropout=None, output_emb_size=None):
         self.output_emb_size = output_emb_size
         if output_emb_size > 0:
             weight_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.TruncatedNormal(std=0.02))
-            self.emb_reduce_linear = paddle.nn.Linear(768, output_emb_size, weight_attr=weight_attr)
+            self.emb_reduce_linear = paddle.nn.Linear(
+                self.ptm.config.hidden_size, output_emb_size, weight_attr=weight_attr
+            )
 
     @paddle.jit.to_static(
         input_spec=[

diff --git a/applications/neural_search/recall/simcse/model.py b/applications/neural_search/recall/simcse/model.py
@@ -12,10 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import abc
-import sys
-
-import numpy as np
 
 import paddle
 import paddle.nn as nn
@@ -36,7 +32,9 @@ def __init__(self, pretrained_model, dropout=None, margin=0.0, scale=20, output_
         self.output_emb_size = output_emb_size
         if output_emb_size > 0:
             weight_attr = paddle.ParamAttr(initializer=paddle.nn.initializer.TruncatedNormal(std=0.02))
-            self.emb_reduce_linear = paddle.nn.Linear(768, output_emb_size, weight_attr=weight_attr)
+            self.emb_reduce_linear = paddle.nn.Linear(
+                self.ptm.config.hidden_size, output_emb_size, weight_attr=weight_attr
+            )
 
         self.margin = margin
         # Used scaling cosine similarity to ease converge
@@ -55,7 +53,7 @@ def get_pooled_embedding(
         # Note: cls_embedding is poolerd embedding with act tanh
         sequence_output, cls_embedding = self.ptm(input_ids, token_type_ids, position_ids, attention_mask)
 
-        if with_pooler == False:
+        if with_pooler is False:
             cls_embedding = sequence_output[:, 0, :]
 
         if self.output_emb_size > 0: