Use cls_before_pooler for DiffCSE models.

voidism · May 4, 2022 · f724b2c · f724b2c
1 parent b3c69a5
commit f724b2c
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/diffcse/tool.py b/diffcse/tool.py
@@ -37,11 +37,9 @@ def __init__(self, model_name_or_path: str,
 
         if pooler is not None:
             self.pooler = pooler
-        elif "unsup" in model_name_or_path:
-            logger.info("Use `cls_before_pooler` for unsupervised models. If you want to use other pooling policy, specify `pooler` argument.")
-            self.pooler = "cls_before_pooler"
         else:
-            self.pooler = "cls"
+            logger.info("Use `cls_before_pooler` for DiffCSE models. If you want to use other pooling policy, specify `pooler` argument.")
+            self.pooler = "cls_before_pooler"
 
     def encode(self, sentence: Union[str, List[str]], 
                 device: str = None,