use hgf ModelOutput class

capreolus-ir · Aug 16, 2022 · f7bfaaa · f7bfaaa
1 parent a0f89dd
commit f7bfaaa
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/capreolus/reranker/TFCEDRKNRM.py b/capreolus/reranker/TFCEDRKNRM.py
@@ -137,11 +137,8 @@ def call(self, x, **kwargs):
         doc_seg = tf.reshape(doc_seg, [batch_size * self.num_passages, self.maxseqlen])
 
         # get BERT embeddings (including CLS) for each passage
-        # TODO switch to hgf's ModelOutput after bumping tranformers version
         outputs = self.bert(doc_input, attention_mask=doc_mask, token_type_ids=doc_seg)
-        if self.config["pretrained"].startswith("bert-"):
-            outputs = (outputs[0], outputs[2])
-        bert_output, all_layer_output = outputs
+        bert_output, all_layer_output = outputs.last_hidden_state, outputs.hidden_states
 
         #  embeddings to create the CLS feature
         cls = bert_output[:, 0, :]