x-transformers 2.2.9__tar.gz → 2.2.11__tar.gz

{x_transformers-2.2.9 → x_transformers-2.2.11}/PKG-INFO

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

{x_transformers-2.2.9 → x_transformers-2.2.11}/pyproject.toml

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

{x_transformers-2.2.9 → x_transformers-2.2.11}/tests/test_x_transformers.py

@@ -800,6 +800,33 @@ def test_entropy_based_tokenizer(
 
     tokenizer(seq[0]) # able to handle without batch dim
 
+def test_entropy_based_tokenizer_max_token_len():
+    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 = 100,
+        max_token_size = 4
+    )
+
+    seq = torch.randint(0, 20000, (1, 16,))
+    lens = torch.tensor([14])
+
+    token_lengths = tokenizer(seq, lens = lens)
+
+    assert token_lengths.amax().item() <= 4
+
 def test_custom_ff_activation():
 
     model = TransformerWrapper(

{x_transformers-2.2.9 → x_transformers-2.2.11}/x_transformers/entropy_based_tokenizer.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from itertools import zip_longest
 
 import torch
@@ -115,6 +116,9 @@ class EntropyBasedTokenizer(Module):
 
             boundaries = boundaries | sub_seq_boundaries
 
+            if exists(mask):
+                boundaries = boundaries & mask
+
         # number of tokens
 
         num_tokens = boundaries.sum(dim = -1)