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

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,31 +150,39 @@ class BeliefStateWrapper(Module):
 
         # get the encoded suffix token once
 
-        if exists(suffix):
-            if suffix.ndim == 1:
-                suffix = repeat(suffix, 'n -> b n', b = batch)
+        if not decode_backwards:
+            if exists(suffix):
+                if suffix.ndim == 1:
+                    suffix = repeat(suffix, 'n -> b n', b = batch)
 
-            suffix = suffix.flip(1) # reverse autoregressive
+                suffix = suffix.flip(1) # reverse autoregressive
 
-        suffix_sos_tokens = rearrange(self.suffix_token, 'd -> 1 1 d')
+            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_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
-        )
+            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
 
-        # pick out the last embedding for fill in the middle
+            prefix_embed = torch.randint(0, self.num_tokens, (batch, 1), device = device)
 
-        suffix_embed = suffix_embed[:, -1:]
+            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,
                 return_intermediates = True,
                 return_embeddings = True,
@@ -181,12 +191,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]
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.20
+Version: 2.1.21
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.1.20.dist-info → x_transformers-2.1.21.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=22jTxhNIKJuQFU8iRanOMpDdyqT_GiCZ2MAprxz6CGo,9841
+x_transformers/belief_state_wrapper.py,sha256=vQUg5djN8TJVhofhPhtMpMbbFz6d1lGsQOukLZhsa3I,10476
 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.20.dist-info/METADATA,sha256=YU5P-lgqBdEofFNiMZH1YIbgH8FddCS-l4K-n1o2h7o,87571
-x_transformers-2.1.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.20.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.20.dist-info/RECORD,,
+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,,