Upload LUAR

model.py

@@ -44,12 +44,12 @@ class LUAR(PreTrainedModel):
     def mean_pooling(self, token_embeddings, attention_mask):
         """Mean Pooling as described in the SBERT paper.
         """
-        input_mask_expanded = repeat(attention_mask, 'b l -> b l d', d=self.hidden_size).float()
+        input_mask_expanded = repeat(attention_mask, 'b l -> b l d', d=self.hidden_size).type(token_embeddings.type())
         sum_embeddings = reduce(token_embeddings * input_mask_expanded, 'b l d -> b d', 'sum')
         sum_mask = torch.clamp(reduce(input_mask_expanded, 'b l d -> b d', 'sum'), min=1e-9)
         return sum_embeddings / sum_mask
     
-    def get_episode_embeddings(self, input_ids, attention_mask):
+    def get_episode_embeddings(self, input_ids, attention_mask, output_attentions=False):
         """Computes the Author Embedding. 
         """
         B, E, _ = attention_mask.shape
@@ -61,7 +61,8 @@ class LUAR(PreTrainedModel):
             input_ids=input_ids,
             attention_mask=attention_mask,
             return_dict=True,
-            output_hidden_states=True
+            output_hidden_states=True,
+            output_attentions=output_attentions,
         )
 
         # at this point, we're embedding individual "comments"
@@ -74,11 +75,14 @@ class LUAR(PreTrainedModel):
         
         episode_embeddings = self.linear(episode_embeddings)
         
+        if output_attentions:
+            return episode_embeddings, outputs["attentions"]
+
         return episode_embeddings
     
-    def forward(self, input_ids, attention_mask):
+    def forward(self, input_ids, attention_mask, output_attentions=False):
         """Calculates a fixed-length feature vector for a batch of episode samples.
         """
-        output = self.get_episode_embeddings(input_ids, attention_mask)
+        output = self.get_episode_embeddings(input_ids, attention_mask, output_attentions)
 
         return output