Add option to CNN baseline to return representation in addition to cl…

…assifier output.
Machine-Learning-for-Medical-Language · Sep 24, 2024 · 8dfcba0 · 8dfcba0
1 parent b0de26e
commit 8dfcba0
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/cnlpt/BaselineModels.py b/src/cnlpt/BaselineModels.py
@@ -67,6 +67,7 @@ def forward(
         input_ids=None,
         event_tokens=None,
         labels=None,
+        output_hidden_states=False,
         **kwargs,
     ):
         embeddings = self.embed(input_ids)
@@ -103,8 +104,10 @@ def forward(
                 loss += self.loss_fn[self.task_names[task_ind]](
                     task_logits, task_labels.type(torch.LongTensor).to(labels.device)
                 )
-
-        return loss, logits
+        if output_hidden_states:
+            return loss, logits, fc_in
+        else:
+            return loss, logits
 
 
 class LstmSentenceClassifier(nn.Module):
@@ -160,5 +163,4 @@ def forward(
                 loss += self.loss_fn(
                     task_logits, task_labels.type(torch.LongTensor).to(labels.device)
                 )
-
         return loss, logits