x-transformers 2.3.24__tar.gz → 2.3.26__tar.gz

{x_transformers-2.3.24 → x_transformers-2.3.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.3.24
+Version: 2.3.26
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.3.24 → x_transformers-2.3.26}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.3.24"
+version = "2.3.26"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.3.24 → x_transformers-2.3.26}/tests/test_x_transformers.py

@@ -1036,3 +1036,43 @@ def test_autoregressive_wrapper(
     loss = wrapper(x)
 
     loss.backward()
+
+def test_prepend_embed():
+
+    from x_transformers import AutoregressiveWrapper
+
+    model = TransformerWrapper(
+        num_tokens = 256,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 12,
+            heads = 8
+        )
+    )
+
+    model = AutoregressiveWrapper(model)
+
+    x = torch.randint(0, 256, (2, 10))
+    prepend_embeds = torch.randn(2, 3, 512)
+
+    loss = model(x, prepend_embeds = prepend_embeds)
+    loss.backward()
+
+    sample = model.generate(
+        prompts = x[:, :1],
+        seq_len = 100,
+        temperature = 0.,
+        prepend_embeds = prepend_embeds,
+        cache_kv = True,
+    )
+
+    sample_no_cache = model.generate(
+        prompts = x[:, :1],
+        seq_len = 100,
+        temperature = 0.,
+        prepend_embeds = prepend_embeds,
+        cache_kv = False,
+    )
+
+    assert torch.allclose(sample, sample_no_cache)

{x_transformers-2.3.24 → x_transformers-2.3.26}/x_transformers/autoregressive_wrapper.py

@@ -309,7 +309,13 @@ class AutoregressiveWrapper(Module):
 
         return out
 
-    def forward(self, x, return_outputs = False, **kwargs):
+    def forward(
+        self,
+        x,
+        return_outputs = False,
+        prepend_embeds = None,
+        **kwargs
+    ):
         seq, ignore_index, add_attn_z_loss, add_next_embed_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss, self.add_continuous_pred_head
 
         inp, target = x, x[:, 1:]
@@ -328,6 +334,7 @@ class AutoregressiveWrapper(Module):
             return_intermediates = True,
             return_attn_z_loss = add_attn_z_loss,
             return_next_embed_pred = add_next_embed_loss,
+            prepend_embeds = prepend_embeds,
             **kwargs
         )
 
@@ -338,6 +345,14 @@ class AutoregressiveWrapper(Module):
         else:
             logits = out
 
+        # if there are prepended embeds, excise it out
+
+        if exists(prepend_embeds):
+            prepend_len = prepend_embeds.shape[1]
+            logits = logits[:, prepend_len:]
+
+        # take all tokens but the last
+
         logits = logits[:, :-1]
 
         # loss function
@@ -356,7 +371,7 @@ class AutoregressiveWrapper(Module):
             loss = loss + cache.attn_z_loss
 
         if add_next_embed_loss:
-            mask = inp[:, :-1] != ignore_index
+            mask = target != ignore_index
             embed_pred = next_embed_pred[:, :-1]
             cont_targets = init_embeds[:, 1:].detach()