ensure pooler consistency

Signed-off-by: Zhiyuan Chen <[email protected]>
DLS5-Omics · Apr 23, 2024 · 262e5c5 · 262e5c5
1 parent 0b0248f
commit 262e5c5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/multimolecule/models/rnabert/modeling_rnabert.py b/multimolecule/models/rnabert/modeling_rnabert.py
@@ -682,6 +682,8 @@ def __init__(self, config: RnaBertConfig):
         self.activation = nn.Tanh()
 
     def forward(self, hidden_states: Tensor) -> Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
         first_token_tensor = hidden_states[:, 0]
         pooled_output = self.dense(first_token_tensor)
         pooled_output = self.activation(pooled_output)

diff --git a/multimolecule/models/rnamsm/modeling_rnamsm.py b/multimolecule/models/rnamsm/modeling_rnamsm.py
@@ -1260,6 +1260,8 @@ def __init__(self, config: RnaMsmConfig):
         self.activation = nn.Tanh()
 
     def forward(self, hidden_states: Tensor) -> Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
         first_token_tensor = hidden_states[:, 0]
         pooled_output = self.dense(first_token_tensor)
         pooled_output = self.activation(pooled_output)