tiny-recursive-model 0.0.2__tar.gz → 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tiny-recursive-model
3
- Version: 0.0.2
3
+ Version: 0.0.3
4
4
  Summary: Tiny Recursive Model
5
5
  Project-URL: Homepage, https://pypi.org/project/tiny-recursive-model/
6
6
  Project-URL: Repository, https://github.com/lucidrains/tiny-recursive-model
@@ -55,6 +55,53 @@ Official repository is [here](https://github.com/SamsungSAILMontreal/TinyRecursi
55
55
 
56
56
  <img width="300" alt="trm-fig3" src="https://github.com/user-attachments/assets/bfe3dd2a-e859-492a-84d5-faf37339f534" />
57
57
 
58
+ ## Install
59
+
60
+ ```bash
61
+ $ pip install tiny-recursive-model
62
+ ```
63
+
64
+ ## Usage
65
+
66
+ ```python
67
+ import torch
68
+ from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D, Trainer
69
+
70
+ trm = TinyRecursiveModel(
71
+ dim = 16,
72
+ num_tokens = 256,
73
+ network = MLPMixer1D(
74
+ dim = 16,
75
+ depth = 2,
76
+ seq_len = 256
77
+ ),
78
+ )
79
+
80
+ from torch.utils.data import Dataset
81
+ class MockDataset(Dataset):
82
+ def __len__(self):
83
+ return 16
84
+
85
+ def __getitem__(self, idx):
86
+ inp = torch.randint(0, 256, (256,))
87
+ out = torch.randint(0, 256, (256,))
88
+ return inp, out
89
+
90
+ trainer = Trainer(
91
+ trm,
92
+ MockDataset(),
93
+ epochs = 1,
94
+ batch_size = 16,
95
+ cpu = True
96
+ )
97
+
98
+ trainer()
99
+
100
+ pred_answer, exit_indices = trm.predict(torch.randint(0, 256, (1, 256)), halt_prob_thres = 0.1)
101
+
102
+ torch.save(trm.state_dict(), 'saved-trm.pt')
103
+ ```
104
+
58
105
  ## Citations
59
106
 
60
107
  ```bibtex
@@ -9,6 +9,53 @@ Official repository is [here](https://github.com/SamsungSAILMontreal/TinyRecursi
9
9
 
10
10
  <img width="300" alt="trm-fig3" src="https://github.com/user-attachments/assets/bfe3dd2a-e859-492a-84d5-faf37339f534" />
11
11
 
12
+ ## Install
13
+
14
+ ```bash
15
+ $ pip install tiny-recursive-model
16
+ ```
17
+
18
+ ## Usage
19
+
20
+ ```python
21
+ import torch
22
+ from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D, Trainer
23
+
24
+ trm = TinyRecursiveModel(
25
+ dim = 16,
26
+ num_tokens = 256,
27
+ network = MLPMixer1D(
28
+ dim = 16,
29
+ depth = 2,
30
+ seq_len = 256
31
+ ),
32
+ )
33
+
34
+ from torch.utils.data import Dataset
35
+ class MockDataset(Dataset):
36
+ def __len__(self):
37
+ return 16
38
+
39
+ def __getitem__(self, idx):
40
+ inp = torch.randint(0, 256, (256,))
41
+ out = torch.randint(0, 256, (256,))
42
+ return inp, out
43
+
44
+ trainer = Trainer(
45
+ trm,
46
+ MockDataset(),
47
+ epochs = 1,
48
+ batch_size = 16,
49
+ cpu = True
50
+ )
51
+
52
+ trainer()
53
+
54
+ pred_answer, exit_indices = trm.predict(torch.randint(0, 256, (1, 256)), halt_prob_thres = 0.1)
55
+
56
+ torch.save(trm.state_dict(), 'saved-trm.pt')
57
+ ```
58
+
12
59
  ## Citations
13
60
 
14
61
  ```bibtex
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "tiny-recursive-model"
3
- version = "0.0.2"
3
+ version = "0.0.3"
4
4
  description = "Tiny Recursive Model"
5
5
  authors = [
6
6
  { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -5,3 +5,7 @@ from tiny_recursive_model.trm import (
5
5
  from tiny_recursive_model.trainer import (
6
6
  Trainer
7
7
  )
8
+
9
+ from tiny_recursive_model.mlp_mixer_1d import (
10
+ MLPMixer1D
11
+ )