titans-pytorch 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

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/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.dist-info → titans_pytorch-0.1.2.dist-info}/METADATA

@@ -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.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
+titans_pytorch/titans.py,sha256=cGWJHkOYmIeE6X383mZvyjusECBwbplVvK0cfgfhBxg,18634
+titans_pytorch-0.1.2.dist-info/METADATA,sha256=FWq5JIp1WY9dYpzatfGzfkcAGQFk-mEPwxF0wCrbM5w,4684
+titans_pytorch-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.2.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.2.dist-info/RECORD,,

titans_pytorch-0.1.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=I-4oF1JPEmUvhLszEocM5cDgHYNFcNp0Q9nbDYSPFqU,195
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=7PHBCbeB1LhHY5s3zAyYF0L3Mm7CNy4TOBbcpLX6bNE,17686
-titans_pytorch/titans.py,sha256=L3Mu6pDnimD4MNn_832trFEJgXOPjxSdTrB9jiSUSTk,18533
-titans_pytorch-0.1.0.dist-info/METADATA,sha256=LuWDzv-NbGxYKOMThN_WKQWDueyIsOAMSwwiE_BDraI,4595
-titans_pytorch-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.0.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.0.dist-info/RECORD,,