x-transformers 2.1.15__tar.gz → 2.1.17__tar.gz

{x_transformers-2.1.15 → x_transformers-2.1.17}/PKG-INFO

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

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

{x_transformers-2.1.15 → x_transformers-2.1.17}/tests/test_x_transformers.py

@@ -732,8 +732,7 @@ def test_belief_state_wrapper(
 
     seq = torch.randint(0, 20000, (2, 16))
 
-    loss = model(seq, backward = False)
-    loss.backward()
+    loss = model(seq) # backwards happen automatically
 
     suffix = None
     if goal_suffix:

x_transformers-2.1.17/train_belief_state.py

@@ -0,0 +1,133 @@
+from x_transformers import TransformerWrapper, Decoder, BeliefStateWrapper
+from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
+
+import random
+import tqdm
+import gzip
+import numpy as np
+import torch
+import torch.optim as optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 2
+GRADIENT_ACCUMULATE_EVERY = 8
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 500
+GENERATE_LENGTH = 256
+SEQ_LEN = 256
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate GPT-like decoder model for forward and backwards
+
+forward_model = TransformerWrapper(
+    num_tokens = 256,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        rotary_pos_emb = True
+    )
+)
+
+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,
+    backward_decoder = backward_model
+)
+
+model.cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    data = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
+    train_x, valid_x = np.split(data, [int(90e6)])
+    data_train, data_val = torch.from_numpy(train_x), torch.from_numpy(valid_x)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
+        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True))
+val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True))
+
+# optimizer
+
+optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10., desc = 'training'):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss = model(
+            next(train_loader),
+            loss_scale = 1./ GRADIENT_ACCUMULATE_EVERY
+        )
+
+    print(f'training loss: {loss.item()}')
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            loss = model(next(val_loader), return_loss_only = True)
+            print(f'validation loss: {loss.item()}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:-1]
+        prime = decode_tokens(inp)
+        print(f'%s \n\n %s', (prime, '*' * 100))
+
+        sample = model.generate_with_suffix_cond(
+            prompts = inp,
+            seq_len = GENERATE_LENGTH,
+            cache_kv = True
+        )
+
+        output_str = decode_tokens(sample)
+        print(output_str)

{x_transformers-2.1.15 → x_transformers-2.1.17}/x_transformers/belief_state_wrapper.py

@@ -23,7 +23,8 @@ from x_transformers.x_transformers import (
 )
 
 import einx
-from einops import rearrange, repeat
+from einops import rearrange, repeat, pack, unpack
+from einops.layers.torch import Rearrange
 
 # helper functions
 
@@ -55,7 +56,9 @@ class BeliefStateWrapper(Module):
         backward_decoder: TransformerWrapper | None = None,
         train_frac_forward_backward_pairs: float = 1.,
         text_head: Module | None = None,
-        backward_ar_loss_weight: float = 1. # can weigh the training of the backwards decoder differently, perhaps fwd/bwd have a shared backbone etc etc
+        backward_ar_loss_weight: float = 1., # can weigh the training of the backwards decoder differently, perhaps fwd/bwd have a shared backbone etc etc
+        pred_terminal = False,
+        pred_terminal_loss_weight: float = 1.
     ):
         super().__init__()
         backward_decoder = default(backward_decoder, forward_decoder) # if backward decoder not set, use the same transformer, assume it knows how to switch gears based on suffix token
@@ -66,6 +69,8 @@ class BeliefStateWrapper(Module):
         dim = forward_decoder.emb_dim
         num_tokens = forward_decoder.num_tokens
 
+        self.to_forward_logits = nn.Linear(dim, num_tokens, bias = False)
+
         # the suffix token
 
         self.suffix_token = nn.Parameter(torch.zeros(dim))
@@ -82,6 +87,17 @@ class BeliefStateWrapper(Module):
 
         self.text_head = text_head
 
+        # predicting terminal state (when suffix and prefix predict the same token)
+
+        self.to_terminal_logit = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.LeakyReLU(),
+            nn.Linear(dim, 1),
+            Rearrange('... 1 -> ...')
+        ) if pred_terminal else None
+
+        self.pred_terminal_loss_weight = pred_terminal_loss_weight
+
         # the two decoders, one which is causal forward, the other causal backwards
 
         self.forward_decoder = forward_decoder
@@ -112,7 +128,7 @@ class BeliefStateWrapper(Module):
         prompts,
         seq_len,
         temperature = 1.25,
-        cache_kv = True,
+        cache_kv = False,
         suffix: Tensor | None = None, # the goal conditioning
         filter_logits_fn = min_p,
         filter_kwargs = dict(
@@ -122,6 +138,8 @@ class BeliefStateWrapper(Module):
     ):
         max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
 
+        prompts, batch_ps = pack([prompts], '* d')
+
         batch, orig_seq_len = prompts.shape
 
         out = prompts
@@ -183,14 +201,19 @@ class BeliefStateWrapper(Module):
 
             # concat sample
 
-            out = torch.cat((out, sample), dim=-1)
+            out = torch.cat((out, sample), dim = -1)
+
+        out = out[:, orig_seq_len:]
 
-        return out[:, orig_seq_len:]
+        out, = unpack(out, batch_ps, '* n')
+
+        return out
 
     def forward(
         self,
         seq,
-        backward = True
+        return_loss_only = False,
+        loss_scale = 1.
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -230,6 +253,7 @@ class BeliefStateWrapper(Module):
         # f - forward, b - backward, i - indices
 
         fi, bi = fb_pairs.unbind(dim = -1)
+
         valid_mask = (bi - fi) >= 2
 
         fb_pairs = fb_pairs[valid_mask]
@@ -251,8 +275,9 @@ class BeliefStateWrapper(Module):
 
         labels_fi, labels_bi = (fi + 1), bi
 
-        forward_labels, backward_labels = seq[:, fi], seq[:, bi]
-        labels = stack((forward_labels, backward_labels), dim = -1)
+        forward_labels, backward_labels = seq[:, labels_fi], seq[:, labels_bi]
+
+        labels = cat((forward_labels, backward_labels), dim = -1)
 
         # get the forward and backward embedding pairs and feed them through the text head for both forward and backward predictions
 
@@ -265,33 +290,47 @@ class BeliefStateWrapper(Module):
 
         # cross entropy loss
 
-        fb_loss = F.cross_entropy(
+        loss = F.cross_entropy(
             rearrange(logits, 'b n (fb l) -> b l (fb n)', fb = 2),
-            rearrange(labels, 'b n fb -> b (fb n)'),
+            labels,
             reduction = 'none' if self.needs_loss_weight else 'mean'
         )
 
+        # maybe predict terminal
+
+        if exists(self.to_terminal_logit):
+            terminal_logits = self.to_terminal_logit(fb_embeds)
+
+            terminal_labels = ((bi - fi) == 2).float() # distance is exactly 2
+            terminal_labels = repeat(terminal_labels, 'n -> b n', b = batch)
+
+            is_end_loss = F.binary_cross_entropy_with_logits(
+                terminal_logits,
+                terminal_labels
+            )
+
+            loss = (
+                loss +
+                is_end_loss * self.pred_terminal_loss_weight
+            )
+
+        # maybe early return loss
+
+        if return_loss_only:
+            return loss
+
         # maybe loss weighting
 
         if self.needs_loss_weight:
-            fb_loss = rearrange(fb_loss, 'b (fb n) -> b fb n', fb = 2)
-            fb_loss = einx.multiply('b fb n, fb', fb_loss, self.loss_weights)
-            fb_loss = fb_loss.mean()
+            loss = rearrange(loss, 'b (fb n) -> b fb n', fb = 2)
+            loss = einx.multiply('b fb n, fb', loss, self.loss_weights)
+            loss = loss.mean()
 
         # backwards
 
-        orig_backward = getattr(fb_loss, 'backward')
-
-        def patched_backward_fn(*args, **kwargs):
-            orig_backward(*args, **kwargs)
-            orig_forward_embeds.backward(forward_embeds.grad)
-            orig_backward_embeds.backward(backward_embeds.grad)
-
-        # can allow the researcher to call .backward from the outside
+        (loss * loss_scale).backward()
 
-        if backward:
-            patched_backward_fn()
-        else:
-            setattr(fb_loss, 'backward', patched_backward_fn)
+        orig_forward_embeds.backward(forward_embeds.grad)
+        orig_backward_embeds.backward(backward_embeds.grad)
 
-        return fb_loss
+        return loss