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

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.3.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.3
+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
@@ -77,6 +77,8 @@ trm = TinyRecursiveModel(
     ),
 )
 
+# mock dataset
+
 from torch.utils.data import Dataset
 class MockDataset(Dataset):
     def __len__(self):
@@ -87,9 +89,13 @@ class MockDataset(Dataset):
         out = torch.randint(0, 256, (256,))
         return inp, out
 
+mock_dataset = MockDataset()
+
+# trainer
+
 trainer = Trainer(
     trm,
-    MockDataset(),
+    mock_dataset,
     epochs = 1,
     batch_size = 16,
     cpu = True
@@ -97,9 +103,18 @@ trainer = Trainer(
 
 trainer()
 
-pred_answer, exit_indices = trm.predict(torch.randint(0, 256, (1, 256)), halt_prob_thres = 0.1)
+# 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

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

@@ -1,8 +0,0 @@
-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=Ep18uwvhWjHxGeyv42ruXLVc2F6TlZg2_CmeVVfYz7c,6001
-tiny_recursive_model-0.0.3.dist-info/METADATA,sha256=0enBPVOxRoReOf0hms_ZoAI4HHdMWUrrW4Ps0MTuQ9g,3943
-tiny_recursive_model-0.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-tiny_recursive_model-0.0.3.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-tiny_recursive_model-0.0.3.dist-info/RECORD,,