x-transformers 2.2.7__tar.gz → 2.2.9__tar.gz

{x_transformers-2.2.7 → x_transformers-2.2.9}/PKG-INFO

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

{x_transformers-2.2.7 → x_transformers-2.2.9}/pyproject.toml

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

{x_transformers-2.2.7 → x_transformers-2.2.9}/tests/test_x_transformers.py

@@ -798,4 +798,24 @@ def test_entropy_based_tokenizer(
 
     assert len(segmented_seq) == seq.shape[0]
 
-    tokenizer(seq[0]) # able to handle without batch dim
+    tokenizer(seq[0]) # able to handle without batch dim
+
+def test_custom_ff_activation():
+
+    model = TransformerWrapper(
+        num_tokens = 20000,
+        max_seq_len = 1024,
+        attn_layers = Decoder(
+            dim = 128,
+            depth = 6,
+            heads = 8,
+            attn_dim_head = 64,
+            ff_custom_activation = nn.Sigmoid()
+        )
+    )
+
+    seq = torch.randint(0, 20000, (2, 1024))
+
+    logits = model(seq)
+
+    assert logits.shape == (2, 1024, 20000)

x_transformers-2.2.9/train_entropy_tokenizer.py

@@ -0,0 +1,118 @@
+from x_transformers import TransformerWrapper, Decoder
+from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
+from x_transformers.entropy_based_tokenizer import EntropyBasedTokenizer
+
+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 = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 100
+GENERATE_LENGTH = 1024
+SEQ_LEN = 1024
+
+# 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
+
+model = TransformerWrapper(
+    num_tokens = 256,
+    max_seq_len = SEQ_LEN,
+    attn_layers = Decoder(
+        dim = 512,
+        depth = 6,
+        heads = 8,
+        rotary_pos_emb = True
+    )
+)
+
+tokenizer = EntropyBasedTokenizer(
+    model,
+    entropy_threshold = 2.5
+)
+
+model = AutoregressiveWrapper(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 / GRADIENT_ACCUMULATE_EVERY).backward()
+
+    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))
+            print(f'validation loss: {loss.item()}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:-1]
+
+        tokens = tokenizer(inp, return_segmented_seq = True)
+
+        delimiter = " \u275A "
+        output_str = delimiter.join([decode_tokens(token) for token in tokens])
+
+        print(f"{output_str}\n\n")

{x_transformers-2.2.7 → x_transformers-2.2.9}/x_transformers/entropy_based_tokenizer.py

@@ -33,12 +33,15 @@ class EntropyBasedTokenizer(Module):
     def __init__(
         self,
         decoder: Module,
-        entropy_threshold: float
+        entropy_threshold: float,
+        max_token_size: int | None = None
     ):
         super().__init__()
         self.decoder = decoder
         self.entropy_threshold = entropy_threshold
 
+        self.max_token_size = max_token_size
+
     @torch.no_grad()
     def forward(
         self,
@@ -53,7 +56,7 @@ class EntropyBasedTokenizer(Module):
         self.decoder.eval()
 
         is_var_length = exists(lens)
-        batch, seq_len, device = *seq.shape, seq.device
+        batch, seq_len, device, max_token_size = *seq.shape, seq.device, self.max_token_size
 
         arange = torch.arange(seq_len, device = device)
 
@@ -94,7 +97,29 @@ class EntropyBasedTokenizer(Module):
             scatter_indices = rearrange(lens - 1, 'b -> b 1')
             boundaries.scatter_(-1, scatter_indices, True)
 
-        num_tokens = boundaries.sum(dim = -1) # number of tokens
+        # handle max token size - technique has the flaw that repeating subsequences are grouped into one large token
+
+        if exists(max_token_size):
+            token_ids = boundaries.cumsum(dim = -1)
+            token_ids = F.pad(token_ids, (1, -1), value = 0)
+
+            max_num_tokens = boundaries.sum(dim = -1).amax().item()
+            token_ids_seq = torch.arange(max_num_tokens, device = device)
+
+            token_mask = einx.equal('j, b i -> b j i', token_ids_seq, token_ids)
+
+            token_sub_seq_arange = token_mask.cumsum(dim = -1)
+
+            sub_seq_boundaries = (token_sub_seq_arange % max_token_size == 0)
+            sub_seq_boundaries = (sub_seq_boundaries & token_mask).any(dim = 1)
+
+            boundaries = boundaries | sub_seq_boundaries
+
+        # number of tokens
+
+        num_tokens = boundaries.sum(dim = -1)
+
+        # get number of tokens as well as derived indices
 
         indices = arange_plus_one[boundaries].split(num_tokens.tolist())
 

{x_transformers-2.2.7 → x_transformers-2.2.9}/x_transformers/x_transformers.py

@@ -1196,6 +1196,7 @@ class FeedForward(Module):
         glu_mult_bias = False,
         swish = False,
         relu_squared = False,
+        custom_activation = None,
         post_act_ln = False,
         dropout = 0.,
         no_bias = False,
@@ -1205,7 +1206,9 @@ class FeedForward(Module):
         inner_dim = int(dim * mult)
         dim_out = default(dim_out, dim)
 
-        if relu_squared:
+        if exists(custom_activation):
+            activation = deepcopy(custom_activation)
+        elif relu_squared:
             activation = ReluSquared()
         elif swish:
             activation = nn.SiLU()