Upload AbLang

encoderblocks.py

@@ -13,7 +13,7 @@ class AbRepOutput():
     """
     Dataclass used to store AbRep output.
     """
-    last_hidden_states: torch.FloatTensor
+    last_hidden_state: torch.FloatTensor
     all_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
 
@@ -36,7 +36,7 @@ class EncoderBlocks(PreTrainedModel):
                 all_hidden_states = all_hidden_states + (hidden_states,) # Takes out each hidden states after each EncoderBlock
             if output_attentions: 
                 all_self_attentions = all_self_attentions + (attentions,) # Takes out attention layers for analysis
-        return AbRepOutput(last_hidden_states=hidden_states, all_hidden_states=all_hidden_states, attentions=all_self_attentions)
+        return AbRepOutput(last_hidden_state=hidden_states, all_hidden_states=all_hidden_states, attentions=all_self_attentions)
     
     
 class EncoderBlock(PreTrainedModel):

model.py

@@ -47,8 +47,8 @@ def apply_cls_embeddings(inputs, outputs):
     for i in d:
         mask[i, d[i]] = 0
     mask[:, 0] = 0.0 # make cls token invisible
-    mask = mask.unsqueeze(-1).expand(outputs.last_hidden_states.size())
-    sum_embeddings = torch.sum(outputs.last_hidden_states * mask, 1)
+    mask = mask.unsqueeze(-1).expand(outputs.last_hidden_state.size())
+    sum_embeddings = torch.sum(outputs.last_hidden_state * mask, 1)
     sum_mask = torch.clamp(mask.sum(1), min=1e-9)
-    outputs.last_hidden_states[:, 0, :] = sum_embeddings / sum_mask
+    outputs.last_hidden_state[:, 0, :] = sum_embeddings / sum_mask
     return outputs