fix first-last avg

Author: Tianyu Gao

evaluation.py

@@ -125,7 +125,7 @@ def batcher(params, batch, max_length=None):
         elif args.pooler == "avg":
             return ((last_hidden * batch['attention_mask'].unsqueeze(-1)).sum(1) / batch['attention_mask'].sum(-1).unsqueeze(-1)).cpu()
         elif args.pooler == "avg_first_last":
-            first_hidden = hidden_states[0]
+            first_hidden = hidden_states[1]
             last_hidden = hidden_states[-1]
             pooled_result = ((first_hidden + last_hidden) / 2.0 * batch['attention_mask'].unsqueeze(-1)).sum(1) / batch['attention_mask'].sum(-1).unsqueeze(-1)
             return pooled_result.cpu()

simcse/models.py

@@ -70,7 +70,7 @@ def forward(self, attention_mask, outputs):
         elif self.pooler_type == "avg":
             return ((last_hidden * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(-1).unsqueeze(-1))
         elif self.pooler_type == "avg_first_last":
-            first_hidden = hidden_states[0]
+            first_hidden = hidden_states[1]
             last_hidden = hidden_states[-1]
             pooled_result = ((first_hidden + last_hidden) / 2.0 * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(-1).unsqueeze(-1)
             return pooled_result