Upload layers.py

layers.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import enumerate_spans
+from torch.nn.utils.rnn import pad_sequence
+
+class SpanEnumerationLayer(nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+
+    def compute_embeddings(self,embeddings, enumerations, operation = 'sum'):
+    
+      computed_embeddings = []
+
+      for enumeration, embedding in zip(enumerations, embeddings):
+          
+          output_embeddings = []
+          dim_size = embedding.shape[-1]
+          
+          for idx in range(len(enumeration)):
+            
+            x1,x2 = enumeration[idx]
+            output_tensor = embedding[x1:x2+1].sum(dim=0)
+
+            if(operation == 'average'):
+              divisor = abs((x2+1)-x1)
+              output_tensor=torch.div(output_tensor, divisor)  
+
+            output_embeddings.append(output_tensor)
+          computed_embeddings.append(torch.stack(output_embeddings))
+
+      return computed_embeddings
+              
+    def forward(self, embeddings, lengths):
+      
+      enumerations = enumerate_spans(lengths)  
+      computed_embeddings = self.compute_embeddings(embeddings, enumerations)
+
+      return computed_embeddings, enumerations
+