x-transformers 2.4.10__py3-none-any.whl → 2.4.12__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -183,11 +183,12 @@ class AutoregressiveWrapper(Module):
         prompts,
         seq_len,
         beams = 4,
+        return_beams_and_scores = False,
         eos_token = None,
         temperature = 1.,
         stochastic = False,
         prompt_lens: Tensor | None = None,
-        filter_logits_fn: str | Callable = top_k,
+        filter_logits_fn: str | Callable = identity,
         restrict_to_max_seq_len = True,
         filter_kwargs: dict = dict(),
         cache_kv = True,
@@ -285,7 +286,7 @@ class AutoregressiveWrapper(Module):
             out = repeat(out, 'b ... -> (b beams) ...', beams = beams)
             samples = rearrange(samples, 'b beams -> (b beams) 1')
 
-            if should_cache:
+            if should_cache and is_first:
                 modify_cached_kv(cache, lambda t: repeat(t, 'b ... -> (b beams) ...', beams = beams))
 
             # concat sample
@@ -303,14 +304,13 @@ class AutoregressiveWrapper(Module):
 
                 scores = scores[:, :beams]
                 top_beams_indices = sort_indices[:, :beams]
-                top_beams_indices = beams * batch_arange[:, None] + top_beams_indices
+
+                top_beams_indices = curr_num_beams * batch_arange[:, None] + top_beams_indices
 
                 flattened_beam_indices = rearrange(top_beams_indices, 'b beams -> (b beams)')
 
                 out = out[flattened_beam_indices]
 
-                modify_cached_kv(cache, lambda t: t[flattened_beam_indices])
-
             scores = rearrange(scores, 'b beams -> (b beams)')
 
             if not exists(eos_token):
@@ -331,11 +331,17 @@ class AutoregressiveWrapper(Module):
 
         out = rearrange(out, '(b beams) seq -> b beams seq', b = batch)
 
-        out = out[:, 0, orig_seq_len:]
+        out = out[..., orig_seq_len:]
 
-        out, = unpack(out, packed_shape, '* n')
+        out, = unpack(out, packed_shape, '* beams n') # prompt may have no batch dimension
 
-        return out
+        if not return_beams_and_scores:
+            return out[..., 0, :]
+
+        scores = rearrange(scores, '(b beams) -> beams b', b = batch)
+        out = rearrange(out, 'b beams n -> beams b n')
+
+        return out, scores
 
     @torch.no_grad()
     @eval_decorator

{x_transformers-2.4.10.dist-info → x_transformers-2.4.12.dist-info}/METADATA

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

{x_transformers-2.4.10.dist-info → x_transformers-2.4.12.dist-info}/RECORD

@@ -1,6 +1,6 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
 x_transformers/attend.py,sha256=Ax34Rw56xXAWfFPqtZ_B8iKEW2EfQdbVoc9uFjfeNjA,17404
-x_transformers/autoregressive_wrapper.py,sha256=oTGHT5k52gsF5MXvH1WxkJ1Zq3HEmQ4laNM96RMdeiY,17594
+x_transformers/autoregressive_wrapper.py,sha256=y798kS9_VvPOY_5Ilits_64aXNqYvGuilsky1y07ryE,17834
 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=9Fi0HvzpeIJqM6HlAd2M6JqsfjhTN1zEH9iFIimyjS4,117608
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.4.10.dist-info/METADATA,sha256=gbkmhEOCOTsxuZ7zwdya-GF2FkdqDjd9ZlVUcBuRHZ8,90224
-x_transformers-2.4.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.4.10.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.4.10.dist-info/RECORD,,
+x_transformers-2.4.12.dist-info/METADATA,sha256=0o0pofz1ZRwIeFs-e-D9fISslemepdDCQne8FHHeccc,90224
+x_transformers-2.4.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.4.12.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.4.12.dist-info/RECORD,,