x-transformers 2.2.8__py3-none-any.whl → 2.2.10__py3-none-any.whl

x_transformers/entropy_based_tokenizer.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from itertools import zip_longest
 
 import torch
@@ -33,12 +34,15 @@ class EntropyBasedTokenizer(Module):
     def __init__(
         self,
         decoder: Module,
-        entropy_threshold: float
+        entropy_threshold: float,
+        max_token_size: int | None = None
     ):
         super().__init__()
         self.decoder = decoder
         self.entropy_threshold = entropy_threshold
 
+        self.max_token_size = max_token_size
+
     @torch.no_grad()
     def forward(
         self,
@@ -53,7 +57,7 @@ class EntropyBasedTokenizer(Module):
         self.decoder.eval()
 
         is_var_length = exists(lens)
-        batch, seq_len, device = *seq.shape, seq.device
+        batch, seq_len, device, max_token_size = *seq.shape, seq.device, self.max_token_size
 
         arange = torch.arange(seq_len, device = device)
 
@@ -94,7 +98,29 @@ class EntropyBasedTokenizer(Module):
             scatter_indices = rearrange(lens - 1, 'b -> b 1')
             boundaries.scatter_(-1, scatter_indices, True)
 
-        num_tokens = boundaries.sum(dim = -1) # number of tokens
+        # handle max token size - technique has the flaw that repeating subsequences are grouped into one large token
+
+        if exists(max_token_size):
+            token_ids = boundaries.cumsum(dim = -1)
+            token_ids = F.pad(token_ids, (1, -1), value = 0)
+
+            max_num_tokens = boundaries.sum(dim = -1).amax().item()
+            token_ids_seq = torch.arange(max_num_tokens, device = device)
+
+            token_mask = einx.equal('j, b i -> b j i', token_ids_seq, token_ids)
+
+            token_sub_seq_arange = token_mask.cumsum(dim = -1)
+
+            sub_seq_boundaries = (token_sub_seq_arange % max_token_size == 0)
+            sub_seq_boundaries = (sub_seq_boundaries & token_mask).any(dim = 1)
+
+            boundaries = boundaries | sub_seq_boundaries
+
+        # number of tokens
+
+        num_tokens = boundaries.sum(dim = -1)
+
+        # get number of tokens as well as derived indices
 
         indices = arange_plus_one[boundaries].split(num_tokens.tolist())
 

{x_transformers-2.2.8.dist-info → x_transformers-2.2.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.2.8
+Version: 2.2.10
 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.8.dist-info → x_transformers-2.2.10.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=c9dcERmoK454mD3JKlxeU3cNwYVPi8I8bI91xrvVvOA,3938
+x_transformers/entropy_based_tokenizer.py,sha256=ZISAQnbLCQMAbOxPXBbkTDkBl5se4Icvr_HrGMJWElA,4941
 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=m2xiiTafFZiII-QZLCpPerdWbY8O41I6BAYCaaPdXig,111953
 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.8.dist-info/METADATA,sha256=U54PnKoqbbDeOAU2moEtVMOAejOvXnF6BDpakJTHud8,88686
-x_transformers-2.2.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.2.8.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.2.8.dist-info/RECORD,,
+x_transformers-2.2.10.dist-info/METADATA,sha256=cMSCK4d3F2o_MLRH_w8MYMsAhptWQug7fbm_wBhswuY,88687
+x_transformers-2.2.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.2.10.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.2.10.dist-info/RECORD,,