x-transformers 2.1.4__py3-none-any.whl → 2.1.5__py3-none-any.whl

x_transformers/belief_state.py

@@ -35,11 +35,12 @@ class BeliefStateWrapper(Module):
     def __init__(
         self,
         forward_decoder: TransformerWrapper,
-        backward_decoder: TransformerWrapper
+        backward_decoder: TransformerWrapper,
+        train_frac_forward_backward_pairs: float = 1.
     ):
         super().__init__()
         assert forward_decoder.emb_dim == backward_decoder.emb_dim, 'forward and backwards model must have the same embedding dimension'
-        assert forward_decoder.num_tokens == backward_decoder.num_tokens, 'forward and backwards model must have the same embedding dimension'
+        assert forward_decoder.num_tokens == backward_decoder.num_tokens, 'forward and backwards model must have the same number of tokens'
 
         dim = forward_decoder.emb_dim
         num_tokens = forward_decoder.num_tokens
@@ -47,6 +48,7 @@ class BeliefStateWrapper(Module):
         # the suffix token
 
         self.suffix_token = nn.Parameter(torch.zeros(dim))
+        nn.init.normal_(self.suffix_token, std = 0.02)
 
         # the text prediction head, which predicts for the combinations of prefix and suffix the next and previous token for forwards and backward sequences
 
@@ -61,6 +63,13 @@ class BeliefStateWrapper(Module):
         self.forward_decoder = forward_decoder
         self.backward_decoder = backward_decoder
 
+        # what fraction of forward backward pairs to train on
+        # for further memory efficiency
+
+        assert 0 < train_frac_forward_backward_pairs <= 1.
+        self.train_frac_fb_pairs = train_frac_forward_backward_pairs
+        self.needs_subsample_fb_pairs = train_frac_forward_backward_pairs < 1.
+
     def forward(
         self,
         seq
@@ -107,6 +116,17 @@ class BeliefStateWrapper(Module):
 
         fb_pairs = fb_pairs[valid_mask]
 
+        # maybe subsample fb pairs
+
+        if self.needs_subsample_fb_pairs:
+            num_pairs = fb_pairs.shape[0]
+
+            num_subsampled = max(int(num_pairs * self.train_frac_fb_pairs), 1)
+
+            rand_subsampled_indices = torch.randperm(num_pairs, device = device)[:num_subsampled]
+
+            fb_pairs = fb_pairs[rand_subsampled_indices]
+
         # get labels for both
 
         fi, bi = fb_pairs.unbind(dim = -1)

{x_transformers-2.1.4.dist-info → x_transformers-2.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.4
+Version: 2.1.5
 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.4.dist-info → x_transformers-2.1.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
 x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
-x_transformers/belief_state.py,sha256=5E_08m6kvXROYgROazTJFxuUsDPNjIVM3AJxg3CJNmU,3966
+x_transformers/belief_state.py,sha256=GfYDeDqmhldozECgFsJ9zhd6O5NMvdYA5OwueVs8SB4,4742
 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=-80N4sqUr3sR51Ms4wCfc4jhxnPwf0ApNR4xfIsasfQ,110142
 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.4.dist-info/METADATA,sha256=-jme9jyXVeVlo1T7nPi2iGFfoxfjSwdj2D33_dr4yxQ,87570
-x_transformers-2.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.4.dist-info/RECORD,,
+x_transformers-2.1.5.dist-info/METADATA,sha256=-nQpm1eBGBXkLEXGiQK06NaIWa13CEnBujBNfTzvnJ8,87570
+x_transformers-2.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.5.dist-info/RECORD,,