x-transformers 2.1.20__tar.gz → 2.1.22__tar.gz

{x_transformers-2.1.20 → x_transformers-2.1.22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.20
+Version: 2.1.22
 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.20 → x_transformers-2.1.22}/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.20 → x_transformers-2.1.22}/pyproject.toml

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

{x_transformers-2.1.20 → x_transformers-2.1.22}/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,
@@ -121,8 +126,11 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10., desc = 'training'):
         model.eval()
         inp = random.choice(val_dataset)[:-1]
         prime = decode_tokens(inp)
+
         print(f'%s \n\n %s', (prime, '*' * 100))
 
+        print('forwards:\n')
+
         sample = model.generate_with_suffix_cond(
             prompts = inp,
             seq_len = GENERATE_LENGTH,
@@ -131,3 +139,15 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10., desc = 'training'):
 
         output_str = decode_tokens(sample)
         print(output_str)
+
+        print('\nbackwards:\n')
+
+        sample = model.generate_with_suffix_cond(
+            prompts = inp,
+            seq_len = GENERATE_LENGTH,
+            cache_kv = True,
+            decode_backwards = True
+        )
+
+        output_str = decode_tokens(sample.flip(0))
+        print(output_str)

{x_transformers-2.1.20 → x_transformers-2.1.22}/x_transformers/belief_state_wrapper.py

@@ -34,15 +34,6 @@ def exists(v):
 def default(v, d):
     return v if exists(v) else d
 
-def eval_decorator(fn):
-    def inner(self, *args, **kwargs):
-        was_training = self.training
-        self.eval()
-        out = fn(self, *args, **kwargs)
-        self.train(was_training)
-        return out
-    return inner
-
 # wrappers
 
 class BeliefStateWrapper(Module):
@@ -69,6 +60,8 @@ class BeliefStateWrapper(Module):
         dim = forward_decoder.emb_dim
         num_tokens = forward_decoder.num_tokens
 
+        self.num_tokens = num_tokens
+
         # the suffix token
 
         self.suffix_token = nn.Parameter(torch.zeros(dim))
@@ -132,6 +125,7 @@ class BeliefStateWrapper(Module):
         filter_kwargs = dict(
             min_p = 0.1
         ),
+        decode_backwards = False,
         **kwargs
     ):
         max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
@@ -140,6 +134,14 @@ class BeliefStateWrapper(Module):
 
         batch, orig_seq_len = prompts.shape
 
+        # allow for decoding backwards, to make sure it is working
+
+        main_decoder = self.forward_decoder
+
+        if decode_backwards:
+            prompts = prompts.flip(1)
+            main_decoder = self.backward_decoder
+
         out = prompts
 
         # kv caches
@@ -148,32 +150,41 @@ class BeliefStateWrapper(Module):
 
         # get the encoded suffix token once
 
-        if exists(suffix):
-            if suffix.ndim == 1:
-                suffix = repeat(suffix, 'n -> b n', b = batch)
-
-            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,
-            return_embeddings = True
-        )
+        if not decode_backwards:
+            if exists(suffix):
+                if suffix.ndim == 1:
+                    suffix = repeat(suffix, 'n -> b n', b = batch)
 
-        # pick out the last embedding for fill in the middle
+                suffix = suffix.flip(1) # reverse autoregressive
 
-        suffix_embed = suffix_embed[:, -1:]
+            suffix_embed = self.backward_decoder(
+                suffix,
+                prepend_embeds = suffix_sos_tokens,
+                return_embeddings = True
+            )
+
+            # pick out the last embedding for fill in the middle
+
+            suffix_embed = suffix_embed[:, -1:]
+
+        else:
+            # just grab a random token for now for prefix
+
+            prefix_embed = torch.randint(0, self.num_tokens, (batch, 1), device = device)
+
+            prefix_embed = self.forward_decoder(prefix_embed, return_embeddings = True)
 
         # sampling up to seq_len
 
         for _ in range(seq_len):
 
-            embeds, new_cache = self.forward_decoder(
+            embeds, new_cache = main_decoder(
                 out,
+                prepend_embeds = suffix_sos_tokens if decode_backwards else None,
                 return_intermediates = True,
                 return_embeddings = True,
                 cache = cache,
@@ -181,12 +192,18 @@ class BeliefStateWrapper(Module):
             )
 
             last_embeds = embeds[:, -1:]
-            embeds = cat((last_embeds, suffix_embed), dim = -1)
+
+            if not decode_backwards:
+                embeds = cat((last_embeds, suffix_embed), dim = -1)
+            else:
+                embeds = cat((prefix_embed, last_embeds), dim = -1)
 
             if cache_kv and self.forward_decoder.can_cache_kv:
                 cache = new_cache
 
-            logits, _ = self.text_head(embeds).chunk(2, dim = -1)
+            forward_logits, backward_logits = self.text_head(embeds).chunk(2, dim = -1)
+
+            logits = forward_logits if not decode_backwards else backward_logits
 
             logits = logits[:, -1]