x-transformers 2.1.36__tar.gz → 2.2.0__tar.gz

{x_transformers-2.1.36 → x_transformers-2.2.0}/PKG-INFO

@@ -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 → x_transformers-2.2.0}/README.md

@@ -2416,4 +2416,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 → x_transformers-2.2.0}/pyproject.toml

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

{x_transformers-2.1.36 → x_transformers-2.2.0}/tests/test_x_transformers.py

@@ -768,3 +768,25 @@ def test_dynamic_tanh():
     x = torch.randint(0, 20000, (2, 1024))
 
     model(x)
+
+def test_entropy_based_tokenizer():
+    from x_transformers.entropy_based_tokenizer import EntropyBasedTokenizer
+
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 128,
+            depth = 6,
+            heads = 8,
+            attn_dim_head = 64,
+        )
+    )
+
+    tokenizer = EntropyBasedTokenizer(model, entropy_threshold = 9.738)
+
+    seq = torch.randint(0, 20000, (2, 1024))
+
+    segmented_seq = tokenizer(seq, return_segmented_seq = True)
+
+    assert len(segmented_seq) == seq.shape[0]

{x_transformers-2.1.36 → x_transformers-2.2.0}/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-2.2.0/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-2.1.36 → x_transformers-2.2.0}/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: