x-transformers 1.32.12__py3-none-any.whl → 1.32.15__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -19,6 +19,9 @@ def default(val, d):
 def identity(t, *args, **kwargs):
     return t
 
+def join(arr, delimiter = ', '):
+    return delimiter.join(arr)
+
 def cast_tuple(t, length = 1):
     return t if isinstance(t, tuple) else (t,) * length
 
@@ -79,11 +82,29 @@ def top_k(logits, frac_num_tokens = 0.1, k = None):
 # top_a
 
 def top_a(logits, min_p_pow = 2.0, min_p_ratio = 0.02):
-    probs = F.softmax(logits, dim = -1)
-    max_probs = torch.amax(probs, dim = -1, keepdim = True)
+    probs = logits.softmax(dim = -1)
+    max_probs = probs.amax(dim = -1, keepdim = True)
     limit = torch.pow(max_probs, min_p_pow) * min_p_ratio
     return torch.where(probs < limit, float('-inf'), logits)
 
+# min_p
+# https://arxiv.org/abs/2407.01082
+
+def min_p(logits, min_p = 0.1):
+    probs = logits.softmax(dim = -1)
+    max_probs = probs.amax(dim = -1, keepdim = True)
+    limit = min_p * max_probs
+    return torch.where(probs < limit, float('-inf'), logits)
+
+# filter logits functions dict[str -> Callable]
+
+FILTER_LOGITS_FN = dict(
+    top_p = top_p,
+    top_k = top_k,
+    top_a = top_a,
+    min_p = min_p
+)
+
 # contrastive decoding function
 
 def contrastive_decode_fn(
@@ -136,7 +157,7 @@ class AutoregressiveWrapper(Module):
         eos_token = None,
         temperature = 1.,
         prompt_lens: Tensor | None = None,
-        filter_logits_fn: Callable = top_k,
+        filter_logits_fn: str | Callable = top_k,
         restrict_to_max_seq_len = True,
         amateur_model: Module | Tuple[Module] | None = None,
         filter_kwargs: dict = dict(),
@@ -153,6 +174,13 @@ class AutoregressiveWrapper(Module):
 
         b, t = prompts.shape
 
+        # handle filter logits fn given as string
+
+        if isinstance(filter_logits_fn, str):
+            assert filter_logits_fn in FILTER_LOGITS_FN, f"only {join(FILTER_LOGITS_FN.keys())} are available"
+
+            filter_logits_fn = FILTER_LOGITS_FN[filter_logits_fn]
+
         # handle variable lengthed prompts (prefixes)
 
         seq_start_pos = None

x_transformers/x_transformers.py

@@ -1925,7 +1925,8 @@ class TransformerWrapper(Module):
         attn_z_loss_weight = 1e-4,
         average_pool_embed = False,
         use_cls_token = False,
-        squeeze_out_last_dim = False
+        squeeze_out_last_dim = False,
+        token_emb: TokenEmbedding | None = None,
     ):
         super().__init__()
 
@@ -1939,7 +1940,11 @@ class TransformerWrapper(Module):
         self.shift_mem_down = shift_mem_down
 
         self.l2norm_embed = l2norm_embed
-        self.token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed = l2norm_embed)
+
+        if not exists(token_emb):
+            token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed = l2norm_embed)
+
+        self.token_emb = token_emb
 
         no_abs_pos_emb = max_seq_len == 0 or not (use_abs_pos_emb and not attn_layers.disable_abs_pos_emb)
 

{x_transformers-1.32.12.dist-info → x_transformers-1.32.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.32.12
+Version: 1.32.15
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.32.12.dist-info → x_transformers-1.32.15.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
 x_transformers/attend.py,sha256=MI-m91wumBFqFqr_KK9MLgsLk_vPeaVbFMyDr_mWdmY,11349
-x_transformers/autoregressive_wrapper.py,sha256=uX8Mb0zLsQrZECt_9UGt35g7tC05Rk3nPqO6xp2FFCc,9619
+x_transformers/autoregressive_wrapper.py,sha256=ka_iiej5lEBOcbutWQgGrFVMDilz2PFWzLhBh5_tmmg,10366
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=nsuYDfF4GY4kTImXEFqygnpw5mO8DOqaD_PJaeOxFS4,78549
+x_transformers/x_transformers.py,sha256=pyRQ6lb1Sx1CbjOH882tAv9UhAzsLwIeXDPBOsiRipg,78669
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.32.12.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.32.12.dist-info/METADATA,sha256=oOwIIjHp8Bl1ClFKTGaiNAX3RNK46C6jmriZEbyWYvM,662
-x_transformers-1.32.12.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-x_transformers-1.32.12.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.32.12.dist-info/RECORD,,
+x_transformers-1.32.15.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.32.15.dist-info/METADATA,sha256=cXOxhpryz1YtL2rhp8siArEI-K3t9-NOyh-4z1GfjDM,662
+x_transformers-1.32.15.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
+x_transformers-1.32.15.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.32.15.dist-info/RECORD,,

{x_transformers-1.32.12.dist-info → x_transformers-1.32.15.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any