allow variable lengthed contexts

Author: lucidrains

README.md

@@ -244,9 +244,12 @@ seq = torch.randint(0, 4, (1, 196_608 // 2,)).cuda()
 target = torch.randn(1, 200, 4).cuda()      # 4 tracks
 context = torch.randn(4, 16, 1024).cuda()   # 4 contexts for the different 'tracks', each with 16 tokens
 
+context_mask = torch.ones(4, 16).bool().cuda() # optional context mask, in example, include all context tokens
+
 loss = model(
     seq,
     context = context,
+    context_mask = context_mask,
     target = target
 )
 

enformer_pytorch/finetune.py

@@ -149,6 +149,7 @@ def forward(
         seq,
         *,
         context,
+        context_mask = None,
         target = None,
         freeze_enformer = False
     ):
@@ -183,6 +184,13 @@ def forward(
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
         sim = einsum('b h i d, c h j d -> b c h i j', q, k) * self.scale
 
+        # masking
+
+        if exists(context_mask):
+            context_mask = F.pad(context_mask, (1, 0), value = True)
+            context_mask =rearrange(context_mask, 'b j -> b 1 1 1 j')
+            sim = sim.masked_fill(~context_mask, -torch.finfo(sim.dtype).max)
+
         # attention
 
         attn = sim.softmax(dim = -1)

setup.py

@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.1.15',
+  version = '0.1.16',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',