tiny-recursive-model 0.0.2__py3-none-any.whl → 0.0.3__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-0.0.2.dist-info → tiny_recursive_model-0.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tiny-recursive-model
-Version: 0.0.2
+Version: 0.0.3
 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,53 @@ 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
+    ),
+)
+
+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
+
+trainer = Trainer(
+    trm,
+    MockDataset(),
+    epochs = 1,
+    batch_size = 16,
+    cpu = True
+)
+
+trainer()
+
+pred_answer, exit_indices = trm.predict(torch.randint(0, 256, (1, 256)), halt_prob_thres = 0.1)
+
+torch.save(trm.state_dict(), 'saved-trm.pt')
+```
+
 ## Citations
 
 ```bibtex

tiny_recursive_model-0.0.3.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=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,,

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,,