x-transformers 2.5.3__py3-none-any.whl → 2.5.5__py3-none-any.whl

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
 

{x_transformers-2.5.3.dist-info → x_transformers-2.5.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.5.3
+Version: 2.5.5
 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.dist-info → x_transformers-2.5.5.dist-info}/RECORD

@@ -1,6 +1,6 @@
 x_transformers/__init__.py,sha256=aVuhUU0572TJHW88BVc4yA2tla0Zb8l3NH7W4RZ1AEs,1005
 x_transformers/attend.py,sha256=Ax34Rw56xXAWfFPqtZ_B8iKEW2EfQdbVoc9uFjfeNjA,17404
-x_transformers/autoregressive_wrapper.py,sha256=y798kS9_VvPOY_5Ilits_64aXNqYvGuilsky1y07ryE,17834
+x_transformers/autoregressive_wrapper.py,sha256=BsGO9xfVYkvynqbU1__tu_S_cxl7gss0YwnkhIa2baY,18401
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
@@ -12,7 +12,7 @@ x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,
 x_transformers/x_transformers.py,sha256=fW-AoomNCw4n2JFbZN9rZV3lKQvz_Tl6L4txUvac_9o,119993
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.5.3.dist-info/METADATA,sha256=iR77ECuqz3O70zaZ5Mx3NwbNNal-FerMlKPlXTbv8vE,90223
-x_transformers-2.5.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.5.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.5.3.dist-info/RECORD,,
+x_transformers-2.5.5.dist-info/METADATA,sha256=Igay1acyeLzF_vDvB9BW7NWuAy_ck7G2rhITKre3Lew,90223
+x_transformers-2.5.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.5.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.5.5.dist-info/RECORD,,