x-transformers 2.5.3__tar.gz → 2.5.4__tar.gz

{x_transformers-2.5.3 → x_transformers-2.5.4}/PKG-INFO

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

{x_transformers-2.5.3 → x_transformers-2.5.4}/pyproject.toml

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

{x_transformers-2.5.3 → x_transformers-2.5.4}/tests/test_x_transformers.py

@@ -1181,3 +1181,32 @@ def test_attn_pooler(
     out = model(x)
 
     assert out.shape == (2, num_pooled_tokens, 77)
+
+def test_prompts_given_as_list_tensor():
+    from x_transformers import AutoregressiveWrapper
+
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 12,
+            heads = 8
+        )
+    )
+
+    wrapped = AutoregressiveWrapper(model)
+
+    seq = torch.randint(0, 20000, (3, 1024))
+
+    loss = wrapped(seq)
+    loss.backward()
+
+    sampled = wrapped.generate([
+        torch.randint(0, 20000, (3,)),
+        torch.randint(0, 20000, (5,)),
+        torch.randint(0, 20000, (2,)),
+        torch.randint(0, 20000, (7,)),
+    ], 256)
+
+    assert sampled.shape == (4, 256)

{x_transformers-2.5.3 → x_transformers-2.5.4}/x_transformers/autoregressive_wrapper.py

@@ -4,9 +4,10 @@ from math import ceil, log
 from typing import Tuple, Callable
 
 import torch
-from torch import nn, Tensor
+from torch import nn, tensor, Tensor
 from torch.nn import Module
 import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
 
 from einops import rearrange, repeat, pack, unpack
 
@@ -347,7 +348,7 @@ class AutoregressiveWrapper(Module):
     @eval_decorator
     def generate(
         self,
-        prompts,
+        prompts: list[Tensor] | Tensor,
         seq_len,
         eos_token = None,
         temperature = 1.,
@@ -363,11 +364,23 @@ class AutoregressiveWrapper(Module):
         cache_kv = True,
         **kwargs
     ):
-        max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
+        max_seq_len, greedy = self.max_seq_len, temperature == 0.
+
+        # handle prompts given as list of variable lengthed token ids
+
+        if isinstance(prompts, list):
+            assert len(prompts) > 0, 'prompts cannot be empty list'
+            assert not exists(prompt_lens), '`prompt_len` will be auto derived if prompts are passed in as list of Tensors'
+
+            prompt_lens = tensor([t.shape[0] for t in prompts], device = prompts[0].device)
+
+            prompts = pad_sequence(prompts, batch_first = True)
+
+        # pack maybe no batch
 
         prompts, ps = pack([prompts], '* n')
 
-        b, t = prompts.shape
+        b, t, device = *prompts.shape, prompts.device
 
         # handle filter logits fn given as string
 
@@ -380,6 +393,7 @@ class AutoregressiveWrapper(Module):
 
         seq_start_pos = None
         if exists(prompt_lens):
+            print('prompt lens')
             prompts = align_right(prompts, prompt_lens, pad_id = self.pad_value)
             seq_start_pos = t - prompt_lens