x-transformers 2.1.36__py3-none-any.whl → 2.2.0__py3-none-any.whl

x_transformers/__init__.py

@@ -37,3 +37,5 @@ from x_transformers.dpo import (
 from x_transformers.neo_mlp import (
     NeoMLP
 )
+
+from x_transformers.entropy_based_tokenizer import EntropyBasedTokenizer

x_transformers/entropy_based_tokenizer.py

@@ -0,0 +1,91 @@
+import torch
+import torch.nn.functional as F
+from torch.nn import Module
+from torch.nn.utils.rnn import pad_sequence
+
+from x_transformers.x_transformers import Decoder, TransformerWrapper
+
+from einops import repeat, rearrange
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# entropy based tokenizer applied in byte-latent transformer paper
+# they use a simple entropy threshold for segmenting a string into variable sized tokens
+
+# https://arxiv.org/abs/2412.09871
+
+class EntropyBasedTokenizer(Module):
+    def __init__(
+        self,
+        decoder: TransformerWrapper,
+        entropy_threshold = 1.5
+    ):
+        super().__init__()
+        assert isinstance(decoder.attn_layers, Decoder)
+
+        self.decoder = decoder
+        self.entropy_threshold = entropy_threshold
+
+    @torch.no_grad()
+    def forward(
+        self,
+        seq,
+        return_segmented_seq = False
+    ):
+        self.decoder.eval()
+
+        batch, seq_len, device = *seq.shape, seq.device
+
+        _, intermediates = self.decoder(seq, return_logit_entropies = True)
+
+        entropies = intermediates.logit_entropies
+
+        over_thres_mask = entropies >= self.entropy_threshold
+
+        arange = torch.arange(seq_len, device = device) + 1
+        arange = repeat(arange, 'n -> b n', b = batch)
+
+        # get a tensor of Int['b num_tokens'] with the token lengths, zero padded
+
+        boundaries = over_thres_mask.clone()
+        boundaries[..., -1] = True # last token is always a boundary
+
+        num_tokens = boundaries.sum(dim = -1) # number of tokens
+
+        boundaries = arange[boundaries].split(num_tokens.tolist())
+
+        # get the token lengths
+
+        token_lengths = []
+
+        for one_boundary in boundaries:
+            padded_boundary = F.pad(one_boundary, (1, 0), value = 0.)
+            one_token_lengths = padded_boundary[1:] - padded_boundary[:-1]
+
+            token_lengths.append(one_token_lengths)
+
+        token_lengths = pad_sequence(token_lengths, batch_first = True)
+
+        # early return
+
+        if not return_segmented_seq:
+            return token_lengths
+
+        # segment the sequence based on the token lengths
+
+        segmented_seq = []
+
+        for one_seq, one_token_length in zip(seq, token_lengths):
+
+            one_token_length = one_token_length[one_token_length > 0]
+
+            splitted_seq = one_seq.split(one_token_length.tolist())
+            segmented_seq.append(splitted_seq)
+
+        return segmented_seq

x_transformers/x_transformers.py

@@ -2909,6 +2909,7 @@ class TransformerWrapper(Module):
         return_embeddings = False,
         return_logits_and_embeddings = False,
         return_intermediates = False,
+        return_embeddings_and_intermediates = False,
         return_logit_entropies = False,
         mask = None,
         return_mems = False,
@@ -2940,8 +2941,8 @@ class TransformerWrapper(Module):
 
         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
 
-        return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss
-        return_embeddings = return_embeddings | (not exists(self.to_logits))
+        return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss | return_embeddings_and_intermediates
+        return_embeddings = return_embeddings | (not exists(self.to_logits)) | return_embeddings_and_intermediates
 
         # absolute positional embedding
 
@@ -3131,6 +3132,8 @@ class TransformerWrapper(Module):
 
         if return_logits_and_embeddings:
             out = (logits, x)
+        elif return_embeddings_and_intermediates:
+            out = (x, intermediates)
         elif return_embeddings:
             out = x
         else:

{x_transformers-2.1.36.dist-info → x_transformers-2.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.36
+Version: 2.2.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2464,4 +2464,15 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@article{Pagnoni2024ByteLT,
+    title   = {Byte Latent Transformer: Patches Scale Better Than Tokens},
+    author  = {Artidoro Pagnoni and Ram Pasunuru and Pedro Rodriguez and John Nguyen and Benjamin Muller and Margaret Li and Chunting Zhou and Lili Yu and Jason Weston and Luke S. Zettlemoyer and Gargi Ghosh and Mike Lewis and Ari Holtzman and Srinivasan Iyer},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2412.09871},
+    url     = {https://api.semanticscholar.org/CorpusID:274762821}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.1.36.dist-info → x_transformers-2.2.0.dist-info}/RECORD

@@ -1,16 +1,17 @@
-x_transformers/__init__.py,sha256=NDoiBivau559WQ0FvXG4ssU3Il9aoHmTIUFN_1juz0s,911
+x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=p0sCAiH1na236ygwgL1Yyhu36eZBf9cZvoW1JyP_fFE,7073
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
+x_transformers/entropy_based_tokenizer.py,sha256=s56Mfok-ulYTKiYmn06feB0QU91y4SQl1Pgj7W5EO3o,2518
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=voN-uEBEKxpUu9K4MVcneSTrzdgJWnZGuQ1QRZQw4Q4,111596
+x_transformers/x_transformers.py,sha256=twoqq2kfVWxntitHKLs2sxFMhK1CPLxGGBDAmkiHXcM,111812
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.36.dist-info/METADATA,sha256=D0qdMRucK3PWwEi8WwdiJdZ8X_hGTm1r3_7bJzYiWSM,88161
-x_transformers-2.1.36.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.36.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.36.dist-info/RECORD,,
+x_transformers-2.2.0.dist-info/METADATA,sha256=oY5xaR9Xw3prT7RPBw3urS2PKWb7rJADn6SZzRW5Tnw,88686
+x_transformers-2.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.2.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.2.0.dist-info/RECORD,,