x-transformers 2.1.21__tar.gz → 2.1.23__tar.gz

{x_transformers-2.1.21 → x_transformers-2.1.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.21
+Version: 2.1.23
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2444,4 +2444,15 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@article{Charpentier2024GPTOB,
+    title   = {GPT or BERT: why not both?},
+    author  = {Lucas Georges Gabriel Charpentier and David Samuel},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2410.24159},
+    url     = {https://api.semanticscholar.org/CorpusID:273707069}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.1.21 → x_transformers-2.1.23}/README.md

@@ -2396,4 +2396,15 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@article{Charpentier2024GPTOB,
+    title   = {GPT or BERT: why not both?},
+    author  = {Lucas Georges Gabriel Charpentier and David Samuel},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2410.24159},
+    url     = {https://api.semanticscholar.org/CorpusID:273707069}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.1.21 → x_transformers-2.1.23}/pyproject.toml

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

{x_transformers-2.1.21 → x_transformers-2.1.23}/train_belief_state.py

@@ -21,6 +21,8 @@ GENERATE_EVERY  = 500
 GENERATE_LENGTH = 256
 SEQ_LEN = 256
 
+FORWARD_BACKWARD_SAME_MODEL = True
+
 # helpers
 
 def cycle(loader):
@@ -47,16 +49,19 @@ forward_model = TransformerWrapper(
     )
 )
 
-backward_model = TransformerWrapper(
-    num_tokens = 256,
-    max_seq_len = SEQ_LEN,
-    attn_layers = Decoder(
-        dim = 512,
-        depth = 4, # do a smaller backwards
-        heads = 8,
-        rotary_pos_emb = True
+backward_model = None
+
+if not FORWARD_BACKWARD_SAME_MODEL:
+    backward_model = TransformerWrapper(
+        num_tokens = 256,
+        max_seq_len = SEQ_LEN,
+        attn_layers = Decoder(
+            dim = 512,
+            depth = 4, # do a smaller backwards
+            heads = 8,
+            rotary_pos_emb = True
+        )
     )
-)
 
 model = BeliefStateWrapper(
     forward_decoder = forward_model,

{x_transformers-2.1.21 → x_transformers-2.1.23}/x_transformers/belief_state_wrapper.py

@@ -150,6 +150,10 @@ class BeliefStateWrapper(Module):
 
         # get the encoded suffix token once
 
+        suffix_sos_tokens = rearrange(self.suffix_token, 'd -> 1 1 d')
+
+        suffix_sos_tokens = repeat(suffix_sos_tokens, '1 1 d -> b 1 d', b = batch)
+
         if not decode_backwards:
             if exists(suffix):
                 if suffix.ndim == 1:
@@ -157,10 +161,6 @@ class BeliefStateWrapper(Module):
 
                 suffix = suffix.flip(1) # reverse autoregressive
 
-            suffix_sos_tokens = rearrange(self.suffix_token, 'd -> 1 1 d')
-
-            suffix_sos_tokens = repeat(suffix_sos_tokens, '1 1 d -> b 1 d', b = batch)
-
             suffix_embed = self.backward_decoder(
                 suffix,
                 prepend_embeds = suffix_sos_tokens,
@@ -184,6 +184,7 @@ class BeliefStateWrapper(Module):
 
             embeds, new_cache = main_decoder(
                 out,
+                prepend_embeds = suffix_sos_tokens if decode_backwards else None,
                 return_intermediates = True,
                 return_embeddings = True,
                 cache = cache,
@@ -227,7 +228,8 @@ class BeliefStateWrapper(Module):
         self,
         seq,
         return_loss_only = False,
-        loss_scale = 1.
+        loss_scale = 1.,
+        loss_weight_by_fb_indices: callable | None = None
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -335,9 +337,21 @@ class BeliefStateWrapper(Module):
 
         # maybe loss weighting
 
-        if self.needs_loss_weight:
+        needs_loss_weight = default(self.needs_loss_weight, exists(loss_weight_by_fb_indices))
+
+        if needs_loss_weight:
             loss = rearrange(loss, 'b (fb n) -> b fb n', fb = 2)
-            loss = einx.multiply('b fb n, fb', loss, self.loss_weights)
+
+            if self.needs_loss_weight:
+                loss = einx.multiply('b fb n, fb', loss, self.loss_weights)
+
+            # allow researcher to pass in a function that acts on the the forward backward indices Int['n fb']
+            # the reason this may be needed is because the earlier tokens will have more eligible pairs for training, and perhaps this could be normalized
+
+            if exists(loss_weight_by_fb_indices):
+                loss_weight = loss_weight_by_fb_indices(fb_pairs)
+                loss = einx.multiply('b fb n, n', loss, loss_weight)
+
             loss = loss.mean()
 
         # backwards