x-transformers 2.2.3__py3-none-any.whl → 2.2.5__py3-none-any.whl

x_transformers/entropy_based_tokenizer.py

@@ -6,10 +6,14 @@ 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
+from x_transformers.x_transformers import (
+    Decoder,
+    TransformerWrapper,
+    calc_entropy
+)
 
 import einx
-from einops import repeat, rearrange
+from einops import repeat, rearrange, pack, unpack
 
 # helper functions
 
@@ -27,22 +31,23 @@ def default(v, d):
 class EntropyBasedTokenizer(Module):
     def __init__(
         self,
-        decoder: TransformerWrapper,
+        decoder: Module,
         entropy_threshold: float
     ):
         super().__init__()
-        assert isinstance(decoder.attn_layers, Decoder)
-
         self.decoder = decoder
         self.entropy_threshold = entropy_threshold
 
     @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)
@@ -52,9 +57,9 @@ class EntropyBasedTokenizer(Module):
 
         # forward through a small trained decoder and get the entropies of the logits
 
-        _, intermediates = self.decoder(seq, return_logit_entropies = True)
+        logits = self.decoder(seq)
 
-        entropies = intermediates.logit_entropies
+        entropies = calc_entropy(logits)
 
         # get length mask for boundaries
 
@@ -106,6 +111,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
@@ -123,4 +130,7 @@ class EntropyBasedTokenizer(Module):
             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.3.dist-info → x_transformers-2.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.2.3
+Version: 2.2.5
 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.3.dist-info → x_transformers-2.2.5.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=sEOf_J_9PGNFKPZ9Gks3MwNjiTa_JljjO1OU3ubIziQ,3562
+x_transformers/entropy_based_tokenizer.py,sha256=Dr7Hvb0WbdhCeD3_61jAQeONr6IWFFSCRrqBfn0d46Y,3772
 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.3.dist-info/METADATA,sha256=aHA3vvgyUsxooOXzLLPtw0FmRGYrqdzMXe2QVoRJk04,88686
-x_transformers-2.2.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.2.3.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.2.3.dist-info/RECORD,,
+x_transformers-2.2.5.dist-info/METADATA,sha256=5hJR7TPpRznsPZrYAHsbk3Eq9E9XKrAuwybmc3r6B1s,88686
+x_transformers-2.2.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.2.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.2.5.dist-info/RECORD,,