x-transformers 2.2.2__py3-none-any.whl → 2.2.4__py3-none-any.whl

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

{x_transformers-2.2.2.dist-info → x_transformers-2.2.4.dist-info}/METADATA

@@ -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.dist-info → x_transformers-2.2.4.dist-info}/RECORD

@@ -4,14 +4,14 @@ x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67K
 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=hdYfw8GqMj0YVWY_gpaCCzhkMALnQB9yAUaCg8RWMss,3624
+x_transformers/entropy_based_tokenizer.py,sha256=Ac73rJaYL1_lW9VnIZOh7zwJ-n0S96EEOj27pwGMbts,3861
 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=Fl2CuAKTxJDOQvqwQo2FK8eO2s1iNLO-P1PP2Yw64rQ,111826
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.2.2.dist-info/METADATA,sha256=V0g-qeMS5RoayZAttx8bGdkIh0ZAsVmDDfTiPNI5qHM,88686
-x_transformers-2.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.2.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.2.2.dist-info/RECORD,,
+x_transformers-2.2.4.dist-info/METADATA,sha256=N6Sc1vbIHKZXcLjCYRevlCAsbPE71L63Cg68N6gcYxU,88686
+x_transformers-2.2.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.2.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.2.4.dist-info/RECORD,,