x-transformers 2.2.2__tar.gz → 2.2.4__tar.gz

{x_transformers-2.2.2 → x_transformers-2.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.2.2
+Version: 2.2.4
 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.2 → x_transformers-2.2.4}/pyproject.toml

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

{x_transformers-2.2.2 → x_transformers-2.2.4}/tests/test_x_transformers.py

@@ -797,3 +797,5 @@ def test_entropy_based_tokenizer(
     segmented_seq = tokenizer(seq, lens, return_segmented_seq = True)
 
     assert len(segmented_seq) == seq.shape[0]
+
+    tokenizer(seq[0]) # able to handle without batch dim

{x_transformers-2.2.2 → x_transformers-2.2.4}/x_transformers/entropy_based_tokenizer.py

@@ -9,7 +9,7 @@ from torch.nn.utils.rnn import pad_sequence
 from x_transformers.x_transformers import Decoder, TransformerWrapper
 
 import einx
-from einops import repeat, rearrange
+from einops import repeat, rearrange, pack, unpack
 
 # helper functions
 
@@ -39,10 +39,13 @@ class EntropyBasedTokenizer(Module):
     @torch.no_grad()
     def forward(
         self,
-        seq,
-        lens = None, # Int['b']
+        seq,            # Float['b n'] | Float['n']
+        lens = None,    # Int['b']
         return_segmented_seq = False
     ):
+        no_batch_dim = seq.ndim == 1
+        seq, maybe_batch_ps = pack((seq,), '* n')
+
         self.decoder.eval()
 
         is_var_length = exists(lens)
@@ -89,15 +92,15 @@ class EntropyBasedTokenizer(Module):
 
         num_tokens = boundaries.sum(dim = -1) # number of tokens
 
-        boundaries = arange_plus_one[boundaries].split(num_tokens.tolist())
+        indices = arange_plus_one[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]
+        for one_indices in indices:
+            padded_indices = F.pad(one_indices, (1, 0), value = 0.)
+            one_token_lengths = padded_indices[1:] - padded_indices[:-1]
 
             token_lengths.append(one_token_lengths)
 
@@ -106,6 +109,8 @@ class EntropyBasedTokenizer(Module):
         # early return
 
         if not return_segmented_seq:
+            token_lengths, = unpack(token_lengths, maybe_batch_ps, '* num_tokens')
+
             return token_lengths
 
         # segment the sequence based on the token lengths
@@ -120,8 +125,10 @@ class EntropyBasedTokenizer(Module):
 
             one_token_length = one_token_length[one_token_length > 0]
 
-            print(one_seq.shape, one_token_length)
             splitted_seq = one_seq.split(one_token_length.tolist())
             segmented_seq.append(splitted_seq)
 
+        if no_batch_dim:
+            segmented_seq = segmented_seq[0]
+
         return segmented_seq