x-transformers 1.27.6__py3-none-any.whl → 1.27.8__py3-none-any.whl

x_transformers/__init__.py

@@ -23,3 +23,7 @@ from x_transformers.xval import (
 )
 
 from x_transformers.xl_autoregressive_wrapper import XLAutoregressiveWrapper
+
+from x_transformers.dpo import (
+    DPO
+)

x_transformers/continuous.py

@@ -84,6 +84,7 @@ class ContinuousTransformerWrapper(nn.Module):
         mask = None,
         return_attn = False,
         mems = None,
+        mem_masks = None,
         pos = None,
         prepend_embeds = None,
         prepend_mask = None,
@@ -125,7 +126,7 @@ class ContinuousTransformerWrapper(nn.Module):
 
         # attention layers
 
-        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, return_hiddens = True, **kwargs)
+        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, return_hiddens = True, **kwargs)
 
         # splice out memory tokens
 

x_transformers/dpo.py

@@ -0,0 +1,81 @@
+from copy import deepcopy
+
+import torch
+from torch.nn import Module
+import torch.nn.functional as F
+from x_transformers.x_transformers import TransformerWrapper
+
+from einops import rearrange
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def freeze_all_layers_(module):
+    for param in module.parameters():
+        param.requires_grad = False
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def log_prob(prob, indices, eps = 1e-20):
+    indices = rearrange(indices, '... -> ... 1')
+    log_probs = log(prob.gather(-1, indices), eps = eps)
+    return rearrange(log_probs, '... 1 -> ...')
+
+def log_prob_from_model_and_seq(model, seq):
+    logits = model(seq)
+    prob = logits.softmax(dim = -1)
+    return log_prob(prob, seq)
+
+# main class
+
+class DPO(Module):
+    def __init__(
+        self,
+        model: TransformerWrapper,
+        *,
+        beta = 0.1
+    ):
+        super().__init__()
+        self.policy_model = model
+
+        self.ref_model = deepcopy(model)
+        freeze_all_layers_(self.ref_model)
+
+        self.beta = beta
+
+    def parameters(self):
+        return self.policy_model.parameters()
+
+    def forward(
+        self,
+        preferred_seq,
+        unpreferred_seq,
+        prompt_mask = None
+    ):
+        assert preferred_seq.ndim == 2
+        assert preferred_seq.shape == unpreferred_seq.shape
+
+        """
+        Following Appendix B in https://arxiv.org/abs/2305.18290
+        """
+
+        with torch.no_grad():
+            self.ref_model.eval()
+            ref_preferred_logprob = log_prob_from_model_and_seq(self.ref_model, preferred_seq)
+            ref_unpreferred_logprob = log_prob_from_model_and_seq(self.ref_model, unpreferred_seq)
+
+        policy_preferred_logprob = log_prob_from_model_and_seq(self.policy_model, preferred_seq)
+        policy_unpreferred_logprob = log_prob_from_model_and_seq(self.policy_model, unpreferred_seq)
+
+        policy_logratios = policy_preferred_logprob - policy_unpreferred_logprob
+        ref_logratios = ref_preferred_logprob - ref_unpreferred_logprob
+
+        losses = -F.logsigmoid(self.beta * (policy_logratios - ref_logratios))
+
+        if exists(prompt_mask):
+            losses = losses[~prompt_mask]
+
+        return losses.mean()

x_transformers/x_transformers.py

@@ -788,8 +788,8 @@ class Attention(nn.Module):
         # add memory key / values
         self.num_mem_kv = num_mem_kv
         if num_mem_kv > 0:
-            self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
-            self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+            self.mem_k = nn.Parameter(torch.randn(kv_heads, num_mem_kv, dim_head))
+            self.mem_v = nn.Parameter(torch.randn(kv_heads, num_mem_kv, dim_head))
 
         # attention on attention
         self.attn_on_attn = on_attn

{x_transformers-1.27.6.dist-info → x_transformers-1.27.8.dist-info}/METADATA

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

x_transformers-1.27.8.dist-info/RECORD

@@ -0,0 +1,14 @@
+x_transformers/__init__.py,sha256=0-2m0LtLpZiZYGwO-6OMYXofx5hbFb_FJOHMxIBqQr4,673
+x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,10188
+x_transformers/autoregressive_wrapper.py,sha256=gYKIN5Rm8dMYSTX5yHpg9sPYyZf9rsRTJCNrYRdJ-Ww,9618
+x_transformers/continuous.py,sha256=92Wczoaz6dJalix-e3mdIzW0xyRIx3GlBSgsSQOsJeI,6123
+x_transformers/dpo.py,sha256=ek9dgiSs05xeCn8ORceOgKy6LJOnNDw-OJDqxAVLecM,2243
+x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
+x_transformers/x_transformers.py,sha256=c8axLT-n2zz3mvQ1tBbE4KUs-8qL7yFsgtIujyh1JDg,63408
+x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
+x_transformers/xval.py,sha256=ulEPep6i5Hl7H-H9vGfdsmHdprUmK8ajB306jViyV2c,8147
+x_transformers-1.27.8.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.27.8.dist-info/METADATA,sha256=LYDYUsXQOHYBZRr_5pepdN9HSzaW-2nFX5pEzEOFkcA,661
+x_transformers-1.27.8.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+x_transformers-1.27.8.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.27.8.dist-info/RECORD,,

x_transformers-1.27.6.dist-info/RECORD

@@ -1,13 +0,0 @@
-x_transformers/__init__.py,sha256=pXc_U4M3ONUQcpNgZySDIlCF1rp7u4FFmcOYjc4WuXw,629
-x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,10188
-x_transformers/autoregressive_wrapper.py,sha256=gYKIN5Rm8dMYSTX5yHpg9sPYyZf9rsRTJCNrYRdJ-Ww,9618
-x_transformers/continuous.py,sha256=SAZGR-3BgXU7OEQtjg1_9FnrUBpIyVfXfpMrH-oL5rU,6074
-x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=8nBe_MQLfQDHf59pM_c2IiTVHA9frFbFyg4n8S00ZVI,63402
-x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers/xval.py,sha256=ulEPep6i5Hl7H-H9vGfdsmHdprUmK8ajB306jViyV2c,8147
-x_transformers-1.27.6.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.27.6.dist-info/METADATA,sha256=32gbNOf9pJgUoNTdaplhW1mcB4ECiJo5OTCpIVNFWCA,661
-x_transformers-1.27.6.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-x_transformers-1.27.6.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.27.6.dist-info/RECORD,,