x-transformers 2.2.5__tar.gz → 2.2.7__tar.gz

{x_transformers-2.2.5 → x_transformers-2.2.7}/PKG-INFO

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

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

{x_transformers-2.2.5 → x_transformers-2.2.7}/x_transformers/entropy_based_tokenizer.py

@@ -6,12 +6,6 @@ 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,
-    calc_entropy
-)
-
 import einx
 from einops import repeat, rearrange, pack, unpack
 
@@ -23,6 +17,13 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
+def log(t, eps = 1e-20):
+    return t.clamp(min = eps).log()
+
+def calc_entropy_from_logits(logits):
+    prob = logits.softmax(dim = -1)
+    return -(prob * log(prob)).sum(dim = -1)
+
 # entropy based tokenizer applied in byte-latent transformer paper
 # they use a simple entropy threshold for segmenting a string into variable sized tokens
 
@@ -43,7 +44,8 @@ class EntropyBasedTokenizer(Module):
         self,
         seq,            # Float['b n'] | Float['n']
         lens = None,    # Int['b']
-        return_segmented_seq = False
+        return_segmented_seq = False,
+        decoder_forward_kwargs: dict = dict()
     ):
         no_batch_dim = seq.ndim == 1
         seq, maybe_batch_ps = pack((seq,), '* n')
@@ -57,9 +59,9 @@ class EntropyBasedTokenizer(Module):
 
         # forward through a small trained decoder and get the entropies of the logits
 
-        logits = self.decoder(seq)
+        logits = self.decoder(seq, **decoder_forward_kwargs)
 
-        entropies = calc_entropy(logits)
+        entropies = calc_entropy_from_logits(logits)
 
         # get length mask for boundaries