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

x_transformers/entropy_based_tokenizer.py

@@ -1,10 +1,14 @@
+from itertools import zip_longest
+
 import torch
+from torch import tensor
 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
 
+import einx
 from einops import repeat, rearrange
 
 # helper functions
@@ -36,29 +40,56 @@ class EntropyBasedTokenizer(Module):
     def forward(
         self,
         seq,
+        lens = None, # Int['b']
         return_segmented_seq = False
     ):
         self.decoder.eval()
 
+        is_var_length = exists(lens)
         batch, seq_len, device = *seq.shape, seq.device
 
+        arange = torch.arange(seq_len, device = device)
+
+        # forward through a small trained decoder and get the entropies of the logits
+
         _, intermediates = self.decoder(seq, return_logit_entropies = True)
 
         entropies = intermediates.logit_entropies
 
-        over_thres_mask = entropies >= self.entropy_threshold
+        # get length mask for boundaries
+
+        mask = tensor(True, device = device)
+
+        if is_var_length:
+            mask = einx.less('n, b -> b n', arange, lens)
+
+        # the mask for tokens that were of a sufficient surprise level
 
-        arange = torch.arange(seq_len, device = device) + 1
-        arange = repeat(arange, 'n -> b n', b = batch)
+        over_thres_mask = (entropies >= self.entropy_threshold) & mask
+
+        # needed for selecting out indices at entropy threshold mask
+
+        arange_plus_one = arange + 1
+        arange_plus_one = repeat(arange_plus_one, '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
+
+        # set the boundary of the last token
+
+        # if `lens` not given, assume always last token
+        # but if `lens` were given, then properly set the index
+
+        if not is_var_length:
+            boundaries[..., -1] = True
+        else:
+            scatter_indices = rearrange(lens - 1, 'b -> b 1')
+            boundaries.scatter_(-1, scatter_indices, True)
 
         num_tokens = boundaries.sum(dim = -1) # number of tokens
 
-        boundaries = arange[boundaries].split(num_tokens.tolist())
+        boundaries = arange_plus_one[boundaries].split(num_tokens.tolist())
 
         # get the token lengths
 
@@ -79,12 +110,17 @@ class EntropyBasedTokenizer(Module):
 
         # segment the sequence based on the token lengths
 
+        lens = default(lens, (None,))
         segmented_seq = []
 
-        for one_seq, one_token_length in zip(seq, token_lengths):
+        for one_seq, one_len, one_token_length in zip_longest(seq, lens, token_lengths):
+
+            if exists(one_len):
+                one_seq = one_seq[:one_len]
 
             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)
 

x_transformers/x_transformers.py

@@ -2006,7 +2006,7 @@ class AttentionLayers(Module):
         assert not (rotary_xpos and not causal), 'rotary xpos is not compatible with bidirectional attention'
         self.rotary_pos_emb = RotaryEmbedding(rotary_emb_dim, use_xpos = rotary_xpos, scale_base = rotary_xpos_scale_base, interpolation_factor = rotary_interpolation_factor, base_rescale_factor = rotary_base_rescale_factor) if rotary_pos_emb else None
 
-        assert at_most_one_of(alibi_pos_bias, rel_pos_bias, data_dependent_alibi), 'you can only choose one of Alibi positional bias, data dependent Alibi (forgetting transformers), or T5 relative positional bias'
+        assert at_most_one_of(alibi_pos_bias, rel_pos_bias, data_dependent_alibi), 'you can only choose one of Alibi positional bias, data dependent Alibi (forgetting transformers), dynamic tanh, or T5 relative positional bias'
         assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
 
         # relative positional bias

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.2.1
+Version: 2.2.2
 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.1.dist-info → x_transformers-2.2.2.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=2Foh9tBUL55Lu0CcgA8kWzyfHYI7DQHBF3zK_WdQt0o,2519
+x_transformers/entropy_based_tokenizer.py,sha256=hdYfw8GqMj0YVWY_gpaCCzhkMALnQB9yAUaCg8RWMss,3624
 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=twoqq2kfVWxntitHKLs2sxFMhK1CPLxGGBDAmkiHXcM,111812
+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.1.dist-info/METADATA,sha256=SEoe84lxCNvvJ8HVRqaLoESrTYRgQ9dc6Y-Po8nhxbg,88686
-x_transformers-2.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.2.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.2.1.dist-info/RECORD,,
+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,,