x-transformers 1.27.6__tar.gz → 1.27.8__tar.gz

{x-transformers-1.27.6/x_transformers.egg-info → x-transformers-1.27.8}/PKG-INFO

@@ -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.6 → x-transformers-1.27.8}/README.md

@@ -2076,4 +2076,15 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@article{Rafailov2023DirectPO,
+    title   = {Direct Preference Optimization: Your Language Model is Secretly a Reward Model},
+    author  = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Stefano Ermon and Christopher D. Manning and Chelsea Finn},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2305.18290},
+    url     = {https://api.semanticscholar.org/CorpusID:258959321}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x-transformers-1.27.6 → x-transformers-1.27.8}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.27.6',
+  version = '1.27.8',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x-transformers-1.27.6 → x-transformers-1.27.8}/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-1.27.6 → x-transformers-1.27.8}/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-1.27.8/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-1.27.6 → x-transformers-1.27.8}/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 → x-transformers-1.27.8/x_transformers.egg-info}/PKG-INFO

@@ -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.6 → x-transformers-1.27.8}/x_transformers.egg-info/SOURCES.txt

@@ -5,6 +5,7 @@ x_transformers/__init__.py
 x_transformers/attend.py
 x_transformers/autoregressive_wrapper.py
 x_transformers/continuous.py
+x_transformers/dpo.py
 x_transformers/nonautoregressive_wrapper.py
 x_transformers/x_transformers.py
 x_transformers/xl_autoregressive_wrapper.py