tiny-recursive-model 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

tiny_recursive_model/__init__.py

@@ -5,3 +5,7 @@ from tiny_recursive_model.trm import (
 from tiny_recursive_model.trainer import (
     Trainer
 )
+
+from tiny_recursive_model.mlp_mixer_1d import (
+    MLPMixer1D
+)

tiny_recursive_model/trm.py

@@ -25,6 +25,9 @@ def default(v, d):
 def is_empty(t):
     return t.numel() == 0
 
+def range_from_one(n):
+    return range(1, n + 1)
+
 # classes
 
 class TinyRecursiveModel(Module):
@@ -98,11 +101,11 @@ class TinyRecursiveModel(Module):
         latents,   # (b n d)
     ):
 
-        for i in range(self.num_refinement_blocks):
+        for step in range_from_one(self.num_refinement_blocks):
 
             # only last round of refinement receives gradients
 
-            is_last = i == (self.num_refinement_blocks - 1)
+            is_last = step == self.num_refinement_blocks
             context = torch.no_grad if not is_last else nullcontext
 
             with context():
@@ -115,7 +118,7 @@ class TinyRecursiveModel(Module):
         self,
         seq,
         halt_prob_thres = 0.5,
-        num_deep_refinement_steps = 12
+        max_deep_refinement_steps = 12
     ):
         batch = seq.shape[0]
 
@@ -131,9 +134,8 @@ class TinyRecursiveModel(Module):
         exited_step_indices = []
         exited_batch_indices = []
 
-        for i in range(num_deep_refinement_steps):
-            step = i + 1
-            is_last = step == num_deep_refinement_steps
+        for step in range_from_one(max_deep_refinement_steps):
+            is_last = step == max_deep_refinement_steps
 
             outputs, latents = self.deep_refinement(inputs, outputs, latents)
 

{tiny_recursive_model-0.0.2.dist-info → tiny_recursive_model-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tiny-recursive-model
-Version: 0.0.2
+Version: 0.0.4
 Summary: Tiny Recursive Model
 Project-URL: Homepage, https://pypi.org/project/tiny-recursive-model/
 Project-URL: Repository, https://github.com/lucidrains/tiny-recursive-model
@@ -55,6 +55,68 @@ Official repository is [here](https://github.com/SamsungSAILMontreal/TinyRecursi
 
 <img width="300" alt="trm-fig3" src="https://github.com/user-attachments/assets/bfe3dd2a-e859-492a-84d5-faf37339f534" />
 
+## Install
+
+```bash
+$ pip install tiny-recursive-model
+```
+
+## Usage
+
+```python
+import torch
+from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D, Trainer
+
+trm = TinyRecursiveModel(
+    dim = 16,
+    num_tokens = 256,
+    network = MLPMixer1D(
+        dim = 16,
+        depth = 2,
+        seq_len = 256
+    ),
+)
+
+# mock dataset
+
+from torch.utils.data import Dataset
+class MockDataset(Dataset):
+    def __len__(self):
+        return 16
+
+    def __getitem__(self, idx):
+        inp = torch.randint(0, 256, (256,))
+        out = torch.randint(0, 256, (256,))
+        return inp, out
+
+mock_dataset = MockDataset()
+
+# trainer
+
+trainer = Trainer(
+    trm,
+    mock_dataset,
+    epochs = 1,
+    batch_size = 16,
+    cpu = True
+)
+
+trainer()
+
+# inference
+
+pred_answer, exit_indices = trm.predict(
+    torch.randint(0, 256, (1, 256)),
+    max_deep_refinement_steps = 12,
+    halt_prob_thres = 0.1
+)
+
+# save to collection of specialized networks for tool call
+
+torch.save(trm.state_dict(), 'saved-trm.pt')
+
+```
+
 ## Citations
 
 ```bibtex

tiny_recursive_model-0.0.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+tiny_recursive_model/__init__.py,sha256=zuMcrofGu7DnvJM2Mb-O3tqBJF5q8L-8X8OTmq7_o5w,189
+tiny_recursive_model/mlp_mixer_1d.py,sha256=6ivDK9dgHdVl1axg2ayifJ7H5QI3hXptHnb6lfNrno0,1398
+tiny_recursive_model/trainer.py,sha256=6dQPmRaQZWI6527OvlOdgHKCFsufkZnjSHClRdHjs20,4218
+tiny_recursive_model/trm.py,sha256=pvhXZjDFoLQR-bc8ZDc_ikklF-s21cY5o6xYAwSDoL8,6048
+tiny_recursive_model-0.0.4.dist-info/METADATA,sha256=dZKb8mPxFPRghjSG7ZCofEI80nSDlAjFVQAB8JvUAi0,4119
+tiny_recursive_model-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tiny_recursive_model-0.0.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+tiny_recursive_model-0.0.4.dist-info/RECORD,,

tiny_recursive_model-0.0.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-tiny_recursive_model/__init__.py,sha256=obuHzL-k9cpbJiwFxopEYBuMQi898C0r45hqgB6x5Yo,123
-tiny_recursive_model/mlp_mixer_1d.py,sha256=6ivDK9dgHdVl1axg2ayifJ7H5QI3hXptHnb6lfNrno0,1398
-tiny_recursive_model/trainer.py,sha256=6dQPmRaQZWI6527OvlOdgHKCFsufkZnjSHClRdHjs20,4218
-tiny_recursive_model/trm.py,sha256=Ep18uwvhWjHxGeyv42ruXLVc2F6TlZg2_CmeVVfYz7c,6001
-tiny_recursive_model-0.0.2.dist-info/METADATA,sha256=JQeCdRnntKCNIMm5LVmSanMV7pavyPjgWyL3fm2LDJ0,3107
-tiny_recursive_model-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-tiny_recursive_model-0.0.2.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-tiny_recursive_model-0.0.2.dist-info/RECORD,,