x-transformers 1.26.0__py3-none-any.whl → 1.26.1__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1,5 +1,6 @@
 import math
 from random import random
+from typing import Dict
 
 import torch
 from torch import nn, einsum, Tensor
@@ -1464,6 +1465,7 @@ class TransformerWrapper(nn.Module):
         num_tokens,
         max_seq_len,
         attn_layers: AttentionLayers,
+        embed_num_tokens: Dict[str, int] = dict(),
         emb_dim = None,
         max_mem_len = 0,
         shift_mem_down = 0,
@@ -1500,7 +1502,16 @@ class TransformerWrapper(nn.Module):
         else:
             self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len, l2norm_embed = l2norm_embed)
 
-        self.emb_frac_gradient = emb_frac_gradient # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
+        # additional embeddings - say type embedding from BERT
+
+        self.embeds = None
+
+        if len(embed_num_tokens) > 0:
+            self.embeds = nn.ModuleDict({f'{name}_embed': nn.Embedding(num_tokens, emb_dim) for name, num_tokens in embed_num_tokens.items()})
+
+        # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
+
+        self.emb_frac_gradient = emb_frac_gradient
 
         self.post_emb_norm = nn.LayerNorm(emb_dim) if post_emb_norm else nn.Identity()
         self.emb_dropout = nn.Dropout(emb_dropout)
@@ -1548,6 +1559,7 @@ class TransformerWrapper(nn.Module):
         pos = None,
         prepend_embeds = None,
         prepend_mask = None,
+        embed_ids: Dict[str, Tensor] = None,
         sum_embeds = None,
         return_attn_z_loss = False,
         attn_z_loss_weight = 1e-4,
@@ -1564,6 +1576,17 @@ class TransformerWrapper(nn.Module):
         pos_emb = self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos) if not external_pos_emb else pos
         x = self.token_emb(x) + pos_emb
 
+        # add additional embeddings
+
+        if exists(self.embeds) and exists(embed_ids):
+            for name, embed_id in embed_ids.items():
+                embed_key = f'{name}_embed'
+
+                assert embed_key in self.embeds
+                embed = self.embeds[embed_key](embed_id)
+
+                x = x + embed
+
         # for summing embeddings passed externally - needs this for self-conditioning in non-autoregressive training
 
         if exists(sum_embeds):

{x_transformers-1.26.0.dist-info → x_transformers-1.26.1.dist-info}/METADATA

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

{x_transformers-1.26.0.dist-info → x_transformers-1.26.1.dist-info}/RECORD

@@ -3,11 +3,11 @@ x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,1018
 x_transformers/autoregressive_wrapper.py,sha256=2mzOq_rl_vevgrxCDncBlVJOAJGS-XGm-iBKJqMjj_c,9041
 x_transformers/continuous.py,sha256=ixfgi2_zpGN03SX_STXFkNYEOAkgwVIxuS53QgDCx-g,6026
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=iJLui66WnGDjfoYgT0p1Ziv0Bmqpy0ffR8RPkV5x2ck,60868
+x_transformers/x_transformers.py,sha256=k-_pU9lICUYnCumwimPv3VaaxjpKOeFKDdKEWgnvslc,61597
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=lS9W_E_RskPQAqVZkPiUzbByoW1Ajsw_phsikA3JDAg,8139
-x_transformers-1.26.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.26.0.dist-info/METADATA,sha256=wUSqgSth94QEUQWmFjTJFCQxUBa2Cl_aS7QWKP-FTcQ,661
-x_transformers-1.26.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-x_transformers-1.26.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.26.0.dist-info/RECORD,,
+x_transformers-1.26.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.26.1.dist-info/METADATA,sha256=mJsZmH-mfNGzd7eGYglFH6GlNzGUTeugQHLcH8oDecU,661
+x_transformers-1.26.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+x_transformers-1.26.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.26.1.dist-info/RECORD,,