x-transformers 2.1.21__py3-none-any.whl → 2.1.23__py3-none-any.whl

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

{x_transformers-2.1.21.dist-info → x_transformers-2.1.23.dist-info}/METADATA

@@ -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.dist-info → x_transformers-2.1.23.dist-info}/RECORD

@@ -1,7 +1,7 @@
 x_transformers/__init__.py,sha256=NDoiBivau559WQ0FvXG4ssU3Il9aoHmTIUFN_1juz0s,911
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
-x_transformers/belief_state_wrapper.py,sha256=vQUg5djN8TJVhofhPhtMpMbbFz6d1lGsQOukLZhsa3I,10476
+x_transformers/belief_state_wrapper.py,sha256=qDtASUlmc-JcUM0u7pALFgjS9aEsis-jgL07gxWnjsg,11198
 x_transformers/continuous.py,sha256=p0sCAiH1na236ygwgL1Yyhu36eZBf9cZvoW1JyP_fFE,7073
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
@@ -10,7 +10,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 x_transformers/x_transformers.py,sha256=fqgtIs6__JpLWMnJa8AY5OW3AJ2GR1B5p-9TsWdiOIU,110425
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.21.dist-info/METADATA,sha256=kTobyUmA8d0yo8NuwHSrpRchyfBMGZZAa2b0Ly3hZec,87571
-x_transformers-2.1.21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.21.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.21.dist-info/RECORD,,
+x_transformers-2.1.23.dist-info/METADATA,sha256=26PtwpKeLHzR8yI_4ezvQvmuvz-P3uqaYcPK_RVNvcU,87875
+x_transformers-2.1.23.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.23.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.23.dist-info/RECORD,,