Upload LUAR

model.py

@@ -146,7 +146,7 @@ class LUAR(PreTrainedModel):
             config.k_bucket_size,
         )
         self.linear = nn.Linear(self.hidden_size, config.embedding_size)
- 
+
     def create_transformer(self):
         """Creates the Transformer backbone.
         """
@@ -163,7 +163,7 @@ class LUAR(PreTrainedModel):
         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, output_attentions=False):
+    def get_episode_embeddings(self, input_ids, attention_mask, output_attentions=False, document_batch_size=0):
         """Computes the Author Embedding. 
         """
         B, E, _ = attention_mask.shape
@@ -171,14 +171,31 @@ class LUAR(PreTrainedModel):
         input_ids = rearrange(input_ids, 'b e l -> (b e) l')
         attention_mask = rearrange(attention_mask, 'b e l -> (b e) l')
 
-        outputs = self.transformer(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            return_dict=True,
-            output_hidden_states=True,
-            output_attentions=output_attentions,
-        )
-
+        if document_batch_size > 0:
+            outputs = {"last_hidden_state": [], "attentions": []}
+            for i in range(0, len(input_ids), document_batch_size):
+                out = self.transformer(
+                    input_ids=input_ids[i:i+document_batch_size],
+                    attention_mask=attention_mask[i:i+document_batch_size],
+                    return_dict=True,
+                    output_hidden_states=False,
+                    output_attentions=output_attentions,
+                )
+                outputs["last_hidden_state"].append(out["last_hidden_state"])
+                if output_attentions:
+                    outputs["attentions"].append(out["attentions"])
+            outputs["last_hidden_state"] = torch.cat(outputs["last_hidden_state"], dim=0)
+            if output_attentions:
+                outputs["attentions"] = tuple([torch.cat([x[i] for x in outputs["attentions"]], dim=0) for i in range(len(outputs["attentions"][0]))])
+        else:
+            outputs = self.transformer(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                return_dict=True,
+                output_hidden_states=False,
+                output_attentions=output_attentions,
+            )
+            
         # at this point, we're embedding individual "comments"
         comment_embeddings = self.mean_pooling(outputs['last_hidden_state'], attention_mask)
         comment_embeddings = rearrange(comment_embeddings, '(b e) l -> b e l', b=B, e=E)
@@ -194,9 +211,9 @@ class LUAR(PreTrainedModel):
 
         return episode_embeddings
     
-    def forward(self, input_ids, attention_mask, output_attentions=False):
+    def forward(self, input_ids, attention_mask, output_attentions=False, document_batch_size=0):
         """Calculates a fixed-length feature vector for a batch of episode samples.
         """
-        output = self.get_episode_embeddings(input_ids, attention_mask, output_attentions)
+        output = self.get_episode_embeddings(input_ids, attention_mask, output_attentions, document_batch_size)
 
         return output