any text without media should not be updated at all during masked cro…

…ss attention
lucidrains · Jun 8, 2022 · 26ca8db · 26ca8db
1 parent 0e1a915
commit 26ca8db
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/flamingo_pytorch/flamingo_pytorch.py b/flamingo_pytorch/flamingo_pytorch.py
@@ -172,6 +172,12 @@ def forward(
         sim = sim - sim.amax(dim = -1, keepdim = True).detach()
         attn = sim.softmax(dim = -1)
 
+        if exists(media_locations) and self.only_attend_immediate_media:
+            # any text without a preceding media needs to have attention zeroed out
+            text_without_media_mask = text_time == 0
+            text_without_media_mask = rearrange(text_without_media_mask, 'b i -> b 1 i 1')
+            attn.masked_fill(text_without_media_mask, 0.)
+
         out = einsum('... i j, ... j d -> ... i d', attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'flamingo-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.16',
+  version = '0.0.17',
   license='MIT',
   description = 'Flamingo - Pytorch',
   author = 'Phil Wang',