titans-pytorch 0.1.0__tar.gz → 0.1.2__tar.gz

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.1.0
+Version: 0.1.2
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -43,9 +43,9 @@ Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch
 Requires-Dist: tensordict
 Requires-Dist: torch>=2.2
+Requires-Dist: tqdm
 Requires-Dist: x-transformers
 Provides-Extra: examples
-Requires-Dist: tqdm; extra == 'examples'
 Requires-Dist: wandb; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
@@ -104,7 +104,12 @@ transformer = MemoryAsContextTransformer(
 
 token_ids = torch.randint(0, 256, (1, 1023))
 
-logits = transformer(token_ids) # (1, 1023, 256)
+loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)
+loss.backward()
+
+# after much training
+
+sampled = transformer.sample(token_ids[:, :4], 512)
 ```
 
 ## Experiments

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/README.md

@@ -51,7 +51,12 @@ transformer = MemoryAsContextTransformer(
 
 token_ids = torch.randint(0, 256, (1, 1023))
 
-logits = transformer(token_ids) # (1, 1023, 256)
+loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)
+loss.backward()
+
+# after much training
+
+sampled = transformer.sample(token_ids[:, :4], 512)
 ```
 
 ## Experiments

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.0"
+version = "0.1.2"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -34,6 +34,7 @@ dependencies = [
     "rotary-embedding-torch",
     "tensordict",
     "torch>=2.2",
+    "tqdm",
     "x-transformers"
 ]
 
@@ -44,7 +45,6 @@ Repository = "https://github.com/lucidrains/titans-pytorch"
 [project.optional-dependencies]
 
 examples = [
-    "tqdm",
     "wandb"
 ]
 

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/tests/test_titans.py

@@ -3,6 +3,10 @@ from torch import nn
 
 import pytest
 from titans_pytorch import NeuralMemory
+from titans_pytorch.mac_transformer import flex_attention, SegmentedAttention
+
+def exists(v):
+    return v is not None
 
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
@@ -46,10 +50,12 @@ def test_titans_attn_memory():
 
     assert seq.shape == retrieved.shape
 
+@pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('num_persist_mem_tokens', (0, 16))
 @pytest.mark.parametrize('num_longterm_mem_tokens', (0, 16))
 @pytest.mark.parametrize('neural_mem_gate_attn_output', (False, True))
 def test_mac(
+    seq_len,
     num_persist_mem_tokens,
     num_longterm_mem_tokens,
     neural_mem_gate_attn_output
@@ -66,7 +72,32 @@ def test_mac(
         neural_mem_gate_attn_output = neural_mem_gate_attn_output
     )
 
-    x = torch.randint(0, 256, (1, 1023))
+    x = torch.randint(0, 256, (1, seq_len))
 
     logits = transformer(x)
-    assert logits.shape == (1, 1023, 256)
+    assert logits.shape == (1, seq_len, 256)
+
+@pytest.mark.parametrize('seq_len', (1023, 17))
+@pytest.mark.parametrize('sliding', (True, False))
+def test_flex(
+    seq_len,
+    sliding
+):
+    if not (torch.cuda.is_available() and exists(flex_attention)):
+        pytest.skip()
+
+    attn = SegmentedAttention(
+        dim = 512,
+        segment_len = 32,
+        num_persist_mem_tokens = 1,
+        num_longterm_mem_tokens = 1,
+        use_flex_attn = True,
+        sliding = sliding
+    ).cuda()
+
+    seq = torch.randn(1, seq_len, 512).cuda()
+
+    out_flex, _ = attn(seq)
+    out_non_flex, _ = attn(seq, disable_flex_attn = True)
+
+    assert torch.allclose(out_flex, out_non_flex, atol = 1e-5)

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/titans_pytorch/mac_transformer.py

@@ -3,6 +3,8 @@ from typing import Callable
 from math import ceil
 from functools import partial
 
+import tqdm
+
 import torch
 from torch import nn, cat
 import torch.nn.functional as F
@@ -88,12 +90,6 @@ def pad_at_dim(t, pad, dim = -1, value = 0.):
 
 def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
     batch, seq_len = seq.shape[:2]
-
-    need_segment = seq_len >= segment_len
-
-    if not need_segment:
-        return seq, identity
-
     next_seq_len_mult = round_up_multiple(seq_len, segment_len)
 
     padding = next_seq_len_mult - seq_len
@@ -116,6 +112,29 @@ def pad_and_segment_with_inverse(seq, segment_len, fold_into_batch = True):
 
     return seq, inverse
 
+# sampling related
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def gumbel_noise(t):
+    noise = torch.rand_like(t)
+    return -log(-log(noise))
+
+def gumbel_sample(t, temperature = 1.):
+    if temperature > 0.:
+        t = t / temperature + gumbel_noise(t)
+    return t.argmax(dim = -1, keepdim = True)
+
+# min_p
+# https://arxiv.org/abs/2407.01082
+
+def min_p_filter(logits, min_p = 0.1):
+    probs = logits.softmax(dim = -1)
+    max_probs = probs.amax(dim = -1, keepdim = True)
+    limit = min_p * max_probs
+    return torch.where(probs < limit, float('-inf'), logits)
+
 # feedforward and attention
 
 class GEGLU(Module):
@@ -500,6 +519,39 @@ class MemoryAsContextTransformer(Module):
         self.segment_len = segment_len
         self.num_persist_mem_tokens = num_persist_mem_tokens
 
+    @torch.no_grad()
+    def sample(
+        self,
+        prompt: Tensor,
+        seq_len: int,
+        temperature = 1.5,
+        filter_fn: Callable = min_p_filter,
+        filter_kwargs: dict = dict(
+            min_p = 0.1,
+        ),
+        show_progress = True
+    ):
+        was_training = self.training
+        self.eval()
+
+        prompt_seq_len, out = prompt.shape[-1], prompt.clone()
+        sample_num_times = max(0, seq_len - prompt_seq_len)
+
+        iter_wrap = tqdm.tqdm if show_progress else identity
+
+        for _ in iter_wrap(range(sample_num_times)):
+            logits = self.forward(out, disable_flex_attn = True)
+            logits = logits[:, -1]
+
+            logits = filter_fn(logits, **filter_kwargs)
+            sample = gumbel_sample(logits, temperature = temperature)
+
+            out = torch.cat((out, sample), dim = -1)
+
+        self.train(was_training)
+
+        return out[..., prompt_seq_len:]
+
     def forward(
         self,
         x,

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/titans_pytorch/titans.py

@@ -592,7 +592,12 @@ class NeuralMemory(Module):
         batch, seq_len = seq.shape[:2]
 
         if seq_len < self.chunk_size:
-            return self.init_empty_memory_embed(batch, seq_len)
+            out = self.init_empty_memory_embed(batch, seq_len)
+
+            if not return_aux_kv_loss:
+                return out
+
+            return out, self.zero
 
         if exists(past_state):
             past_state = tuple(TensorDict(d) for d in past_state)

{titans_pytorch-0.1.0 → titans_pytorch-0.1.2}/train_mac.py

@@ -68,52 +68,6 @@ def decode_token(token):
 def decode_tokens(tokens):
     return ''.join(list(map(decode_token, tokens)))
 
-# sampling helpers
-
-def log(t, eps = 1e-20):
-    return torch.log(t.clamp(min = eps))
-
-def gumbel_noise(t):
-    noise = torch.zeros_like(t).uniform_(0, 1)
-    return -log(-log(noise))
-
-def gumbel_sample(t, temperature = 1., keepdim = True):
-    if temperature <= 0.:
-        return t.argmax(dim = dim, keepdim = keepdim)
-
-    return ((t / max(temperature, 1e-10)) + gumbel_noise(t)).argmax(dim = -1, keepdim = keepdim)
-
-# min_p
-# https://arxiv.org/abs/2407.01082
-
-def min_p_filter(logits, min_p = 0.1):
-    probs = logits.softmax(dim = -1)
-    max_probs = probs.amax(dim = -1, keepdim = True)
-    limit = min_p * max_probs
-    return torch.where(probs < limit, float('-inf'), logits)
-
-def base_decoding(
-    net,
-    prompt: Tensor,
-    seq_len: int,
-    temperature = 1.5,
-    min_p = 1e-1,
-    filter_thres = 0.9,
-):
-    prompt_seq_len, out = prompt.shape[-1], prompt.clone()
-    sample_num_times = max(0, seq_len - prompt_seq_len)
-
-    for _ in tqdm.tqdm(range(sample_num_times)):
-        logits = net(out, disable_flex_attn = True)
-        logits = logits[:, -1]
-
-        logits = min_p_filter(logits, min_p = min_p)
-        sample = gumbel_sample(logits, temperature = temperature)
-
-        out = torch.cat((out, sample), dim = -1)
-
-    return out[..., prompt_seq_len:]
-
 # instantiate memory-as-context transformer
 
 model = MemoryAsContextTransformer(
@@ -197,6 +151,6 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
         prime = decode_tokens(inp)
         print(f'%s \n\n %s', (prime, '*' * 100))
 
-        sample = base_decoding(model, inp[None, ...], GENERATE_LENGTH)
+        sample = model.sample(inp[None, ...], GENERATE_LENGTH)
         output_str = decode_tokens(sample[0])
         print(output_str)

titans_pytorch-0.1.0/assert_flex.py

@@ -1,18 +0,0 @@
-import torch
-from titans_pytorch.mac_transformer import SegmentedAttention
-
-attn = SegmentedAttention(
-    dim = 512,
-    segment_len = 32,
-    num_persist_mem_tokens = 1,
-    num_longterm_mem_tokens = 1,
-    use_flex_attn = True,
-    sliding = False
-).cuda()
-
-seq = torch.randn(1, 1019, 512).cuda()
-
-out_flex, _ = attn(seq)
-out_non_flex, _ = attn(seq, disable_flex_attn = True)
-
-assert torch.allclose(out_flex, out_non_flex, atol = 1e-5)