@genai-fi/nanogpt 0.23.0 → 1.0.0
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.
- package/README.md +78 -281
- package/dist/{DatasetBuilder-C0iJT29K.js → DatasetBuilder-DU1G1OKX.js} +20 -23
- package/dist/{RealDiv-CNsvC4AU.js → RealDiv-CSnvtN2E.js} +20 -20
- package/dist/{Reshape-dnm9bO3B.js → Reshape-BlylqwWy.js} +12 -12
- package/dist/TeachableLLM.d.ts +10 -15
- package/dist/TeachableLLM.js +201 -2
- package/dist/api/responses.d.ts +81 -0
- package/dist/api/responses.js +169 -0
- package/dist/api/training.d.ts +70 -0
- package/dist/api/training.js +205 -0
- package/dist/data/docx.js +9 -3036
- package/dist/data/stream.d.ts +8 -8
- package/dist/data/stream.js +1 -1
- package/dist/data/textLoader.d.ts +1 -1
- package/dist/data/textLoader.js +2 -2
- package/dist/data.d.ts +3 -0
- package/dist/data.js +12 -0
- package/dist/{dist-BqAU9-yi.js → dist-CwK5S7Ls.js} +2168 -2168
- package/dist/{gpgpu_math-DBYEAAdI.js → gpgpu_math-20tPK8LM.js} +458 -458
- package/dist/{Generator.d.ts → inference/Generator.d.ts} +16 -44
- package/dist/inference/Generator.js +271 -0
- package/dist/inference/tokenisePrompt.d.ts +4 -0
- package/dist/inference/tokenisePrompt.js +13 -0
- package/dist/inference/types.d.ts +44 -8
- package/dist/inference/utilities.d.ts +9 -0
- package/dist/inference/utilities.js +20 -0
- package/dist/jszip.min-DKa1Rjyn.js +3033 -0
- package/dist/{kernel_funcs_utils-D-mATnGR.js → kernel_funcs_utils-ql8Y8qPn.js} +96 -93
- package/dist/layers/MLP.d.ts +1 -1
- package/dist/layers/PositionEmbedding.d.ts +2 -1
- package/dist/layers/PositionEmbedding.js +1 -1
- package/dist/layers/RMSNorm.d.ts +1 -1
- package/dist/layers/TiedEmbedding.js +1 -1
- package/dist/layers.d.ts +4 -0
- package/dist/layers.js +14 -0
- package/dist/loader/load.js +58 -2
- package/dist/loader/loadHF.d.ts +1 -1
- package/dist/loader/loadHF.js +17 -2
- package/dist/loader/loadTransformers.js +46 -2
- package/dist/loader/newZipLoad.js +25 -2
- package/dist/loader/oldZipLoad.d.ts +1 -1
- package/dist/loader/oldZipLoad.js +37 -2
- package/dist/loader/save.js +75 -2
- package/dist/loader/types.d.ts +3 -3
- package/dist/main.d.ts +34 -43
- package/dist/main.js +12327 -20
- package/dist/{matMulGelu-BAIgQaRx.js → matMulGelu-CBoqTZM7.js} +2 -2
- package/dist/models/NanoGPTV1.js +95 -2
- package/dist/models/NanoGPTV2.js +86 -2
- package/dist/models/factory.js +13 -2
- package/dist/models/model.js +76 -2
- package/dist/models.d.ts +4 -0
- package/dist/models.js +14 -0
- package/dist/ops/dot16.js +1 -1
- package/dist/ops/matMulGelu.js +1 -1
- package/dist/ops/webgl/adamAdjust.js +1 -1
- package/dist/ops/webgl/fusedSoftmax.js +2 -2
- package/dist/ops/webgl/gelu.js +2 -2
- package/dist/ops/webgl/log.js +5 -5
- package/dist/ops/webgl/matMulGelu.js +1 -1
- package/dist/ops/webgl/matMulMul.js +1 -1
- package/dist/{stream-BjdpSNqB.js → stream-BpAwcvHz.js} +565 -561
- package/dist/{tfjs_backend-CydPRQTc.js → tfjs_backend-h5weiy1O.js} +36 -36
- package/dist/tokenise.d.ts +4 -0
- package/dist/tokenise.js +15 -0
- package/dist/tokeniser/CharTokeniser.js +18 -20
- package/dist/tokeniser/bpe.js +18 -22
- package/dist/training/BasicTrainer.d.ts +5 -10
- package/dist/training/BasicTrainer.js +80 -88
- package/dist/training/DatasetBuilder.d.ts +4 -4
- package/dist/training/DatasetBuilder.js +1 -1
- package/dist/training/PreTrainer.js +1 -1
- package/dist/training/SFTTrainer.js +1 -1
- package/dist/training/configure.d.ts +3 -0
- package/dist/training/configure.js +32 -0
- package/dist/training/factory.d.ts +6 -0
- package/dist/training/factory.js +8 -0
- package/dist/training/prepareData.d.ts +22 -0
- package/dist/training/prepareData.js +49 -0
- package/dist/training/tasks/TokenStore.d.ts +2 -1
- package/dist/training/tasks/TokenStore.js +8 -5
- package/dist/training/tasks/tokenStream.d.ts +17 -0
- package/dist/training/tasks/tokenStream.js +46 -0
- package/dist/training/types.d.ts +14 -1
- package/dist/training/validateOptions.d.ts +2 -0
- package/dist/training/validateOptions.js +19 -0
- package/dist/training/validation.js +4 -2
- package/dist/utilities/arrayShape.d.ts +1 -0
- package/dist/utilities/arrayShape.js +8 -0
- package/dist/utilities/random.d.ts +1 -0
- package/dist/utilities/random.js +19 -0
- package/dist/utilities/waitForModel.d.ts +1 -1
- package/dist/v4-BK7K-jy_.js +30 -0
- package/package.json +8 -2
- package/dist/Generator.js +0 -2
- package/dist/Trainer-DBsyWJ4s.js +0 -228
- package/dist/Trainer.d.ts +0 -45
- package/dist/Trainer.js +0 -2
- package/dist/main-BSaDGH7I.js +0 -13274
- package/dist/training/tasks/ConversationTask.d.ts +0 -17
- package/dist/training/tasks/ConversationTask.js +0 -29
- package/dist/training/tasks/PretrainingTask.d.ts +0 -17
- package/dist/training/tasks/PretrainingTask.js +0 -42
- package/dist/training/tasks/StartSentenceTask.d.ts +0 -18
- package/dist/training/tasks/StartSentenceTask.js +0 -45
- package/dist/training/tasks/Task.d.ts +0 -29
- package/dist/training/tasks/Task.js +0 -50
- package/dist/training/tasks/splitter.d.ts +0 -5
- package/dist/training/tasks/splitter.js +0 -18
package/README.md
CHANGED
|
@@ -1,29 +1,14 @@
|
|
|
1
1
|
# GenAI NanoGPT
|
|
2
2
|
|
|
3
|
-
A browser-native implementation of
|
|
3
|
+
A browser-native implementation of small transformer language models using TensorFlow.js. This project is an educational toolkit for creating, training and running compact GPT-style models client-side. It supports model creation, tokenisation, dataset preparation, training, and text generation with a single high-level entrypoint: `TeachableLLM`.
|
|
4
4
|
|
|
5
|
-
Live
|
|
5
|
+
Live demo: https://lm.gen-ai.fi
|
|
6
6
|
|
|
7
|
-
|
|
7
|
+
**Design goals**
|
|
8
8
|
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
-
|
|
12
|
-
- **Loading** pre-trained models from various sources (Hugging Face, local files)
|
|
13
|
-
- **Generating** text efficiently on a wide range of devices
|
|
14
|
-
- **Analyzing** model behavior through attention visualization and embeddings
|
|
15
|
-
- **Optimizing** performance across CPU, WebGL, and WebGPU backends
|
|
16
|
-
|
|
17
|
-
### Key Features
|
|
18
|
-
|
|
19
|
-
- 🚀 **Browser-Native**: No server required - train and run models entirely client-side
|
|
20
|
-
- 📱 **Works on Small Devices**: Train models on iPads, phones, and Chromebooks - no powerful hardware needed
|
|
21
|
-
- 🎯 **Multiple Backends**: Automatic backend selection (CPU, WebGL, WebGPU) for optimal performance
|
|
22
|
-
- 🔧 **Flexible Tokenization**: Support for both character-level and BPE tokenizers
|
|
23
|
-
- 📊 **XAI Support**: Attention score visualization, gradient analysis, and embedding extraction
|
|
24
|
-
- 💾 **Model Persistence**: Save and load models in SafeTensors format
|
|
25
|
-
- ⚡ **Performance Optimizations**: Custom WebGPU kernels, gradient checkpointing, and mixed precision training
|
|
26
|
-
- 🎨 **Real-time Training**: Live training metrics and generation during training
|
|
9
|
+
- Small models suitable for experimentation on laptops and mobile devices
|
|
10
|
+
- Clear, teachable APIs for training and generation
|
|
11
|
+
- Browser-first implementation using CPU, WebGL or WebGPU backends
|
|
27
12
|
|
|
28
13
|
## Installation
|
|
29
14
|
|
|
@@ -31,239 +16,111 @@ GenAI NanoGPT is inspired by [Andrej Karpathy's NanoGPT](https://github.com/karp
|
|
|
31
16
|
npm install @genai-fi/nanogpt
|
|
32
17
|
```
|
|
33
18
|
|
|
34
|
-
##
|
|
19
|
+
## Main concepts
|
|
35
20
|
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
import { TeachableLLM, selectBackend } from '@genai-fi/nanogpt';
|
|
21
|
+
- `TeachableLLM` — primary entrypoint. Create, load, save models; access training and responses APIs.
|
|
22
|
+
- `tokenise` — tokeniser helpers and token store (character and BPE tokenisers).
|
|
23
|
+
- `data` — helpers to load text data and stream conversational inputs.
|
|
40
24
|
|
|
41
|
-
|
|
42
|
-
await selectBackend('webgpu'); // or 'webgl', 'cpu'
|
|
25
|
+
The project export surface is centred around `TeachableLLM` (see `lib/main.ts`). This README focuses on the runtime API you will use in applications.
|
|
43
26
|
|
|
44
|
-
|
|
45
|
-
const model = TeachableLLM.create('char', {
|
|
46
|
-
vocabSize: 200,
|
|
47
|
-
blockSize: 128, // Context window size
|
|
48
|
-
nLayer: 4, // Number of transformer layers
|
|
49
|
-
nHead: 4, // Number of attention heads
|
|
50
|
-
nEmbed: 192, // Embedding dimension
|
|
51
|
-
dropout: 0.1,
|
|
52
|
-
useRope: true, // Use Rotary Position Embeddings
|
|
53
|
-
});
|
|
27
|
+
## Quick examples
|
|
54
28
|
|
|
55
|
-
|
|
56
|
-
const trainingText = [
|
|
57
|
-
'The quick brown fox jumps over the lazy dog.',
|
|
58
|
-
'A journey of a thousand miles begins with a single step.',
|
|
59
|
-
// ... more text
|
|
60
|
-
];
|
|
61
|
-
|
|
62
|
-
// Train the model
|
|
63
|
-
await model.train(trainingText, {
|
|
64
|
-
batchSize: 16,
|
|
65
|
-
learningRate: 3e-4,
|
|
66
|
-
maxSteps: 1000,
|
|
67
|
-
logInterval: 10,
|
|
68
|
-
validationSplit: 0.1,
|
|
69
|
-
});
|
|
70
|
-
|
|
71
|
-
// Generate text
|
|
72
|
-
const output = await model.generateText('Once upon a time', {
|
|
73
|
-
maxLength: 100,
|
|
74
|
-
temperature: 0.8,
|
|
75
|
-
topP: 0.9,
|
|
76
|
-
});
|
|
77
|
-
|
|
78
|
-
console.log(output);
|
|
79
|
-
```
|
|
80
|
-
|
|
81
|
-
### Loading a Pre-trained Model
|
|
29
|
+
Creating a model instance
|
|
82
30
|
|
|
83
31
|
```javascript
|
|
84
|
-
import { TeachableLLM
|
|
85
|
-
|
|
86
|
-
// Load from Hugging Face
|
|
87
|
-
const model = TeachableLLM.loadModel('username/model-name');
|
|
88
|
-
|
|
89
|
-
// Or load from a file
|
|
90
|
-
const fileInput = document.getElementById('fileInput');
|
|
91
|
-
fileInput.addEventListener('change', async (event) => {
|
|
92
|
-
const file = event.target.files[0];
|
|
93
|
-
const model = TeachableLLM.loadModel(file);
|
|
94
|
-
await waitForModel(model);
|
|
95
|
-
|
|
96
|
-
const text = await model.generateText('Hello');
|
|
97
|
-
console.log(text);
|
|
98
|
-
});
|
|
99
|
-
```
|
|
32
|
+
import { TeachableLLM } from '@genai-fi/nanogpt';
|
|
100
33
|
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
const model = TeachableLLM.create('char', config);
|
|
109
|
-
|
|
110
|
-
// Listen for training step updates
|
|
111
|
-
model.on('trainStep', (step, progress) => {
|
|
112
|
-
console.log(`Step ${step.step}/${progress.totalSteps}`);
|
|
113
|
-
console.log(`Loss: ${step.loss.toFixed(4)}`);
|
|
114
|
-
console.log(`Validation Loss: ${step.valLoss?.toFixed(4) || 'N/A'}`);
|
|
115
|
-
console.log(`Progress: ${(progress.progress * 100).toFixed(1)}%`);
|
|
116
|
-
console.log(`Time Remaining: ${progress.timeRemaining}s`);
|
|
117
|
-
|
|
118
|
-
// Update UI progress bar
|
|
119
|
-
updateProgressBar(progress.progress);
|
|
120
|
-
updateLossChart(step.loss, step.valLoss);
|
|
34
|
+
// Create a new model with a char or bpe tokeniser
|
|
35
|
+
const model = TeachableLLM.create('char', {
|
|
36
|
+
vocabSize: 200,
|
|
37
|
+
blockSize: 128,
|
|
38
|
+
nLayer: 4,
|
|
39
|
+
nHead: 4,
|
|
40
|
+
nEmbed: 192,
|
|
121
41
|
});
|
|
122
42
|
|
|
123
|
-
|
|
43
|
+
// Switch backend if needed
|
|
44
|
+
await TeachableLLM.selectBackend('webgpu');
|
|
124
45
|
```
|
|
125
46
|
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
Stream generated tokens as they're produced:
|
|
47
|
+
Training the tokeniser (when using streamed conversational data)
|
|
129
48
|
|
|
130
49
|
```javascript
|
|
131
|
-
|
|
50
|
+
import { data, tokenise } from '@genai-fi/nanogpt';
|
|
132
51
|
|
|
133
|
-
//
|
|
134
|
-
|
|
135
|
-
// tokens is an array of new token IDs
|
|
136
|
-
const text = model.tokeniser.decode(tokens);
|
|
137
|
-
console.log('New tokens:', text);
|
|
52
|
+
// Prepare streams using data.loadTextData or MemoryConversationStream
|
|
53
|
+
const streams = await data.loadTextData(['Some example text', 'More text']);
|
|
138
54
|
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
// Generation lifecycle events
|
|
144
|
-
generator.on('start', () => {
|
|
145
|
-
console.log('Generation started');
|
|
146
|
-
showSpinner();
|
|
147
|
-
});
|
|
148
|
-
|
|
149
|
-
generator.on('stop', () => {
|
|
150
|
-
console.log('Generation complete');
|
|
151
|
-
hideSpinner();
|
|
152
|
-
});
|
|
153
|
-
|
|
154
|
-
generator.on('error', (error) => {
|
|
155
|
-
console.error('Generation error:', error);
|
|
156
|
-
});
|
|
157
|
-
|
|
158
|
-
// Start generation
|
|
159
|
-
await generator.generate('Once upon a time', {
|
|
160
|
-
maxLength: 200,
|
|
161
|
-
temperature: 0.8,
|
|
162
|
-
});
|
|
55
|
+
// Train the tokeniser on streams
|
|
56
|
+
const tokens = await model.trainTokeniser(streams);
|
|
57
|
+
console.log('Trained token count:', tokens);
|
|
163
58
|
```
|
|
164
59
|
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
GenAI NanoGPT is designed to work efficiently on resource-constrained devices like iPads, phones, and Chromebooks:
|
|
168
|
-
|
|
169
|
-
### Recommended Settings for Small Devices
|
|
60
|
+
Start a training job
|
|
170
61
|
|
|
171
62
|
```javascript
|
|
172
|
-
|
|
173
|
-
const mobileModel = TeachableLLM.create('char', {
|
|
174
|
-
vocabSize: 200,
|
|
175
|
-
blockSize: 128, // Smaller context window
|
|
176
|
-
nLayer: 4, // Fewer layers
|
|
177
|
-
nHead: 3, // Fewer attention heads
|
|
178
|
-
nEmbed: 192, // Smaller embeddings
|
|
179
|
-
});
|
|
63
|
+
const job = await model.training.job(options, streams, datasets);
|
|
180
64
|
|
|
181
|
-
//
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
maxSteps: 500,
|
|
186
|
-
validationSplit: 0.1,
|
|
187
|
-
logInterval: 50,
|
|
188
|
-
gradientCheckpointing: true,
|
|
189
|
-
mixedPrecision: true,
|
|
65
|
+
// Listen for training progress
|
|
66
|
+
model.training.on('progress', (job) => {
|
|
67
|
+
console.log('Training job progress:', job.progress);
|
|
68
|
+
console.log('Latest log entry:', job.history?.[job.history.length - 1]);
|
|
190
69
|
});
|
|
191
|
-
```
|
|
192
|
-
|
|
193
|
-
### Tips for Training on Mobile Devices
|
|
194
70
|
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
3. **Use Character Tokenization**: Character-level tokenizers use less memory than BPE
|
|
198
|
-
4. **Optimize Training Data**: Use smaller datasets or train in stages
|
|
199
|
-
|
|
200
|
-
## Advanced Usage
|
|
71
|
+
// Pause, resume, cancel via training API using the returned job id
|
|
72
|
+
```
|
|
201
73
|
|
|
202
|
-
|
|
74
|
+
Generate text (responses API)
|
|
203
75
|
|
|
204
76
|
```javascript
|
|
205
|
-
|
|
77
|
+
// Create a response — returns an id and may stream tokens via callback
|
|
78
|
+
const resp = await model.responses.create(
|
|
79
|
+
{
|
|
80
|
+
input: 'Once upon a time',
|
|
81
|
+
maxLength: 100,
|
|
82
|
+
temperature: 0.9,
|
|
83
|
+
},
|
|
84
|
+
(chunk) => {
|
|
85
|
+
// called for intermediate chunks when provided
|
|
86
|
+
console.log('Partial output:', chunk.output);
|
|
87
|
+
}
|
|
88
|
+
);
|
|
206
89
|
|
|
207
|
-
|
|
208
|
-
attentionScores: true,
|
|
209
|
-
maxLength: 50,
|
|
210
|
-
});
|
|
211
|
-
|
|
212
|
-
// Get attention data for visualization
|
|
213
|
-
const attentionData = generator.getAttentionData();
|
|
214
|
-
// Shape: [num_tokens][num_layers][num_heads][seq_len][seq_len]
|
|
90
|
+
console.log('Final output:', resp.output);
|
|
215
91
|
|
|
216
|
-
|
|
217
|
-
//
|
|
92
|
+
// Manage responses
|
|
93
|
+
// model.responses.cancel(id)
|
|
94
|
+
// model.responses.hook(id)
|
|
95
|
+
// model.responses.resume(id)
|
|
218
96
|
```
|
|
219
97
|
|
|
220
|
-
|
|
98
|
+
Tokenisers and token stores
|
|
221
99
|
|
|
222
100
|
```javascript
|
|
223
|
-
|
|
101
|
+
import { tokenise } from '@genai-fi/nanogpt';
|
|
224
102
|
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
updateDisplay(tokens);
|
|
228
|
-
});
|
|
103
|
+
// Character and BPE tokenisers are available
|
|
104
|
+
const { CharTokeniser, BPETokeniser, TokenStore, createTokenStore } = tokenise;
|
|
229
105
|
|
|
230
|
-
|
|
231
|
-
generator.on('stop', () => console.log('Generation complete'));
|
|
232
|
-
|
|
233
|
-
await generator.generate('Once upon a time', {
|
|
234
|
-
maxLength: 200,
|
|
235
|
-
});
|
|
106
|
+
// Use TokenStore to persist prepared token sequences for training
|
|
236
107
|
```
|
|
237
108
|
|
|
238
|
-
|
|
109
|
+
Data helpers
|
|
239
110
|
|
|
240
111
|
```javascript
|
|
241
|
-
|
|
242
|
-
model.enableProfiler = true;
|
|
112
|
+
import { data } from '@genai-fi/nanogpt';
|
|
243
113
|
|
|
244
|
-
//
|
|
245
|
-
const
|
|
246
|
-
if (profiler) {
|
|
247
|
-
console.log('Memory stats:', profiler.getStats());
|
|
248
|
-
}
|
|
114
|
+
// Load plain text into conversation streams
|
|
115
|
+
const streams = await data.loadTextData(['Line one', 'Line two']);
|
|
249
116
|
|
|
250
|
-
//
|
|
251
|
-
|
|
117
|
+
// MemoryConversationStream is useful for in-memory conversations
|
|
118
|
+
const { MemoryConversationStream } = data;
|
|
252
119
|
```
|
|
253
120
|
|
|
254
|
-
## Examples
|
|
255
|
-
|
|
256
|
-
See the [`browser-tests`](browser-tests/) directory for complete examples:
|
|
257
|
-
|
|
258
|
-
- [`generate.html`](browser-tests/generate.html): Text generation with UI
|
|
259
|
-
- [`rope-train.html`](browser-tests/rope-train.html): Training a model with RoPE
|
|
260
|
-
- [`hf.html`](browser-tests/hf.html): Loading from Hugging Face
|
|
261
|
-
- [`loader.html`](browser-tests/loader.html): Loading different file formats
|
|
262
|
-
- [`perf.html`](browser-tests/perf.html): Performance testing
|
|
263
|
-
|
|
264
121
|
## Development
|
|
265
122
|
|
|
266
|
-
|
|
123
|
+
Clone and install dependencies:
|
|
267
124
|
|
|
268
125
|
```bash
|
|
269
126
|
git clone https://github.com/knicos/genai-nanogpt.git
|
|
@@ -271,84 +128,24 @@ cd genai-nanogpt
|
|
|
271
128
|
npm install
|
|
272
129
|
```
|
|
273
130
|
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
```bash
|
|
277
|
-
npm run build # Build for production
|
|
278
|
-
npm run dev # Development mode with watch
|
|
279
|
-
```
|
|
280
|
-
|
|
281
|
-
### Testing
|
|
282
|
-
|
|
283
|
-
```bash
|
|
284
|
-
npm test # Run all tests
|
|
285
|
-
```
|
|
286
|
-
|
|
287
|
-
### Browser Tests
|
|
288
|
-
|
|
289
|
-
```bash
|
|
290
|
-
npm run test:gl # Start dev server
|
|
291
|
-
```
|
|
292
|
-
|
|
293
|
-
### Project Structure
|
|
294
|
-
|
|
295
|
-
```
|
|
296
|
-
lib/
|
|
297
|
-
├── models/ # Model architectures (NanoGPT)
|
|
298
|
-
├── layers/ # Transformer layers (attention, MLP, etc.)
|
|
299
|
-
├── ops/ # Custom TensorFlow.js operations
|
|
300
|
-
│ ├── cpu/ # CPU kernels
|
|
301
|
-
│ ├── webgl/ # WebGL kernels
|
|
302
|
-
│ └── webgpu/ # WebGPU kernels
|
|
303
|
-
├── training/ # Training utilities and optimizers
|
|
304
|
-
├── tokeniser/ # Tokenization implementations
|
|
305
|
-
├── loader/ # Model loading/saving
|
|
306
|
-
├── utilities/ # Helper functions
|
|
307
|
-
└── TeachableLLM.ts # Main API
|
|
308
|
-
```
|
|
309
|
-
|
|
310
|
-
### Custom Operations
|
|
311
|
-
|
|
312
|
-
This library implements several custom TensorFlow.js operations optimized for transformer models:
|
|
313
|
-
|
|
314
|
-
- **RoPE**: Rotary Position Embeddings
|
|
315
|
-
- **Attention Mask**: Causal attention masking
|
|
316
|
-
- **RMS Norm**: Root Mean Square normalization
|
|
317
|
-
- **Adam Optimizer**: Extended Adam with weight decay
|
|
318
|
-
- **16-bit Operators**: To enable mixed-precision training
|
|
319
|
-
|
|
320
|
-
See [`lib/ops`](lib/ops/) for implementations.
|
|
321
|
-
|
|
322
|
-
### Contributing
|
|
323
|
-
|
|
324
|
-
1. Fork the repository
|
|
325
|
-
2. Create a feature branch: `git checkout -b feature/amazing-feature`
|
|
326
|
-
3. Commit your changes: `git commit -m 'Add amazing feature'`
|
|
327
|
-
4. Push to the branch: `git push origin feature/amazing-feature`
|
|
328
|
-
5. Open a Pull Request
|
|
329
|
-
|
|
330
|
-
### Code Style
|
|
331
|
-
|
|
332
|
-
This project uses ESLint and Prettier for code formatting:
|
|
131
|
+
Build and run browser tests:
|
|
333
132
|
|
|
334
133
|
```bash
|
|
335
|
-
npm run
|
|
134
|
+
npm run build
|
|
135
|
+
npm run dev
|
|
136
|
+
npm test
|
|
137
|
+
npm run test:gl
|
|
336
138
|
```
|
|
337
139
|
|
|
338
|
-
##
|
|
140
|
+
## Examples and demos
|
|
339
141
|
|
|
340
|
-
|
|
341
|
-
2. **Batch Size**: Larger batches improve GPU utilization but require more memory
|
|
342
|
-
3. **Mixed Precision**: Enable for faster training on supported hardware (coming soon)
|
|
343
|
-
4. **Gradient Checkpointing**: Reduce memory usage during training, but slower
|
|
344
|
-
5. **Use RoPE**: More efficient than absolute position embeddings
|
|
345
|
-
6. **Start Small on Mobile**: Use 2-4 layers and batch size 2-8 on phones/tablets
|
|
142
|
+
See the `browser-tests/` directory for small example pages demonstrating generation, training, and model loading.
|
|
346
143
|
|
|
347
144
|
## Acknowledgments
|
|
348
145
|
|
|
349
|
-
-
|
|
350
|
-
-
|
|
351
|
-
-
|
|
146
|
+
- Inspired by Andrej Karpathy's NanoGPT: https://github.com/karpathy/nanoGPT
|
|
147
|
+
- Built with TensorFlow.js: https://www.tensorflow.org/js
|
|
148
|
+
- Developed as part of the Finnish Generation AI research project: https://generation-ai-stn.fi
|
|
352
149
|
|
|
353
150
|
## Citation
|
|
354
151
|
|
|
@@ -58,17 +58,17 @@ function b(e) {
|
|
|
58
58
|
return e != null && !ArrayBuffer.isView(e) && (Array.isArray(e) || typeof e == "object" && !(e instanceof f) && !(e instanceof Promise) && !t);
|
|
59
59
|
}
|
|
60
60
|
function x(e) {
|
|
61
|
-
return e == null ||
|
|
61
|
+
return e == null || ee(e) || Array.isArray(e) || typeof e == "object" && e instanceof f || l(e);
|
|
62
62
|
}
|
|
63
|
-
function
|
|
63
|
+
function ee(e) {
|
|
64
64
|
return e === null || typeof e != "object" && typeof e != "function";
|
|
65
65
|
}
|
|
66
66
|
//#endregion
|
|
67
67
|
//#region node_modules/@tensorflow/tfjs-data/dist/util/deep_clone.js
|
|
68
|
-
function
|
|
69
|
-
return h(e,
|
|
68
|
+
function S(e) {
|
|
69
|
+
return h(e, C);
|
|
70
70
|
}
|
|
71
|
-
function
|
|
71
|
+
function C(e) {
|
|
72
72
|
return e instanceof f ? {
|
|
73
73
|
value: e.clone(),
|
|
74
74
|
recurse: !1
|
|
@@ -249,7 +249,7 @@ var O = class {
|
|
|
249
249
|
};
|
|
250
250
|
let e = this.items[this.trav];
|
|
251
251
|
return this.trav++, {
|
|
252
|
-
value:
|
|
252
|
+
value: S(e),
|
|
253
253
|
done: !1
|
|
254
254
|
};
|
|
255
255
|
}
|
|
@@ -683,9 +683,9 @@ var ne = class {
|
|
|
683
683
|
let { batchSize: n = 32, noShuffle: r = !1, ignoreIndex: i = 65535 } = t || {}, a = e.getTokenCount();
|
|
684
684
|
if (a < this.blockSize + 1) throw Error(`Not enough tokens (${a}) for block size ${this.blockSize}`);
|
|
685
685
|
let o = Math.ceil(e.shardSize / this.blockSize), c = {
|
|
686
|
-
shuffledShards: new
|
|
687
|
-
shuffledIndexes: new
|
|
688
|
-
lastShardIndexes: new
|
|
686
|
+
shuffledShards: new Uint16Array(e.getShardCount()),
|
|
687
|
+
shuffledIndexes: new Uint16Array(o),
|
|
688
|
+
lastShardIndexes: new Uint16Array(Math.ceil(e.getShardLength(e.getShardCount() - 1) / this.blockSize)),
|
|
689
689
|
currentMask: null,
|
|
690
690
|
nextMask: null,
|
|
691
691
|
currentShard: null,
|
|
@@ -703,22 +703,19 @@ var ne = class {
|
|
|
703
703
|
})), {
|
|
704
704
|
dataset: Y(async function* () {
|
|
705
705
|
for (;;) {
|
|
706
|
-
let t = (c.shardIndex === c.shuffledShards.length - 1 ? c.lastShardIndexes : c.shuffledIndexes)[c.step]
|
|
707
|
-
if (!
|
|
708
|
-
|
|
709
|
-
|
|
710
|
-
|
|
711
|
-
|
|
712
|
-
|
|
713
|
-
|
|
714
|
-
let e = 0, t = o.slice(n + 1, n + this.blockSize + 1);
|
|
715
|
-
for (let n = 0; n < u.length; n++) t[n] === 0 && (u[n] = i, e++);
|
|
716
|
-
if (e === u.length) continue;
|
|
706
|
+
let t = (c.shardIndex === c.shuffledShards.length - 1 ? c.lastShardIndexes : c.shuffledIndexes)[c.step] * this.blockSize, n = c.currentShard, a = c.currentMask, o = $(c, e, r);
|
|
707
|
+
if (!n) break;
|
|
708
|
+
t + this.blockSize + 1 > n.length && (t = n.length - this.blockSize - 1);
|
|
709
|
+
let s = new Int32Array(n.slice(t, t + this.blockSize)), l = new Int32Array(n.slice(t + 1, t + this.blockSize + 1));
|
|
710
|
+
if (a) {
|
|
711
|
+
let e = 0, n = a.slice(t + 1, t + this.blockSize + 1);
|
|
712
|
+
for (let t = 0; t < l.length; t++) n[t] === 0 && (l[t] = i, e++);
|
|
713
|
+
if (e === l.length) continue;
|
|
717
714
|
}
|
|
718
715
|
yield {
|
|
719
|
-
xs:
|
|
720
|
-
ys:
|
|
721
|
-
}, await
|
|
716
|
+
xs: s,
|
|
717
|
+
ys: l
|
|
718
|
+
}, await o;
|
|
722
719
|
}
|
|
723
720
|
}.bind(this)).batch(n).map((e) => {
|
|
724
721
|
let t = e;
|
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
import { Dn as e, En as t, Io as n, Ks as r, Ms as i, Si as a, Tn as o, nc as s, oc as c, wn as l, xn as u } from "./dist-Da20xy8E.js";
|
|
2
2
|
import { L as d } from "./backend_util-DAV0I34P.js";
|
|
3
|
-
import {
|
|
4
|
-
import {
|
|
5
|
-
import {
|
|
6
|
-
import {
|
|
3
|
+
import { Z as f } from "./gpgpu_math-20tPK8LM.js";
|
|
4
|
+
import { b as p, n as m } from "./kernel_funcs_utils-ql8Y8qPn.js";
|
|
5
|
+
import { t as h } from "./Reshape-BlylqwWy.js";
|
|
6
|
+
import { J as g, b as _ } from "./shared-CfTzpULd.js";
|
|
7
7
|
//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/mean_gpu.js
|
|
8
8
|
var v = class {
|
|
9
9
|
constructor(e, t) {
|
|
@@ -246,7 +246,7 @@ var w = class {
|
|
|
246
246
|
let n = Array(e.length);
|
|
247
247
|
for (let r = 0; r < n.length; r++) n[r] = e[t[r]];
|
|
248
248
|
if (this.outputShape = n, this.rank = n.length, this.rank > 6) throw Error(`Packed transpose for rank ${this.rank} is not yet supported.`);
|
|
249
|
-
let r = f(this.rank), i =
|
|
249
|
+
let r = f(this.rank), i = p("rc", this.rank), a = Array(this.rank);
|
|
250
250
|
for (let e = 0; e < t.length; e++) a[t[e]] = i[e];
|
|
251
251
|
let o = `vec2(${a.slice(-2).join()})`, s = `++${i[this.rank - 1]} < ${n[this.rank - 1]}`, c = `getChannel(getA(${a.join()}), ${o})`;
|
|
252
252
|
this.userCode = `
|
|
@@ -278,15 +278,15 @@ function T(e, t, n) {
|
|
|
278
278
|
//#endregion
|
|
279
279
|
//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/kernels/Sum_impl.js
|
|
280
280
|
function E(n, r, i, d) {
|
|
281
|
-
let f = r, p = n.shape.length, m = s(f, n.shape),
|
|
282
|
-
v && (y = T(n,
|
|
283
|
-
let [b, S] = l(y.shape,
|
|
281
|
+
let f = r, p = n.shape.length, m = s(f, n.shape), g = m, _ = t(g, p), v = _ != null, y = n;
|
|
282
|
+
v && (y = T(n, _, d), g = e(g.length, p)), u("sum", g, p);
|
|
283
|
+
let [b, S] = l(y.shape, g), C = b;
|
|
284
284
|
i && (C = o(b, m));
|
|
285
|
-
let w = c(S), E = c(n.shape) / w, D =
|
|
285
|
+
let w = c(S), E = c(n.shape) / w, D = h({
|
|
286
286
|
inputs: { x: y },
|
|
287
287
|
attrs: { shape: [E, w] },
|
|
288
288
|
backend: d
|
|
289
|
-
}), O = x(D, a(n.dtype), "sum", d), k =
|
|
289
|
+
}), O = x(D, a(n.dtype), "sum", d), k = h({
|
|
290
290
|
inputs: { x: O },
|
|
291
291
|
attrs: { shape: C },
|
|
292
292
|
backend: d
|
|
@@ -307,11 +307,11 @@ var O = {
|
|
|
307
307
|
//#endregion
|
|
308
308
|
//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/kernels/Max_impl.js
|
|
309
309
|
function k(e, t, n, r) {
|
|
310
|
-
let i = c(t), a = c(e.shape) / i, o =
|
|
310
|
+
let i = c(t), a = c(e.shape) / i, o = h({
|
|
311
311
|
inputs: { x: e },
|
|
312
312
|
attrs: { shape: [a, i] },
|
|
313
313
|
backend: r
|
|
314
|
-
}), s = x(o, e.dtype, "max", r), l =
|
|
314
|
+
}), s = x(o, e.dtype, "max", r), l = h({
|
|
315
315
|
inputs: { x: s },
|
|
316
316
|
attrs: { shape: n },
|
|
317
317
|
backend: r
|
|
@@ -321,24 +321,24 @@ function k(e, t, n, r) {
|
|
|
321
321
|
//#endregion
|
|
322
322
|
//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/kernels/Max.js
|
|
323
323
|
function A(n) {
|
|
324
|
-
let { inputs: r, backend: i, attrs: a } = n, { x: d } = r, { reductionIndices: f, keepDims:
|
|
324
|
+
let { inputs: r, backend: i, attrs: a } = n, { x: d } = r, { reductionIndices: f, keepDims: p } = a, m = d.shape.length, h = s(f, d.shape), v = h, y = t(v, m), b = y != null, x = i.shouldExecuteOnCPU([d]), S = d;
|
|
325
325
|
if (b) {
|
|
326
326
|
if (x) {
|
|
327
|
-
let e = i.texData.get(S.dataId).values, t = Array(
|
|
327
|
+
let e = i.texData.get(S.dataId).values, t = Array(m);
|
|
328
328
|
for (let e = 0; e < t.length; e++) t[e] = d.shape[y[e]];
|
|
329
|
-
let n =
|
|
329
|
+
let n = g(e, d.shape, d.dtype, y, t);
|
|
330
330
|
S = i.makeTensorInfo(t, d.dtype);
|
|
331
331
|
let r = i.texData.get(S.dataId);
|
|
332
332
|
r.values = n;
|
|
333
333
|
} else S = T(d, y, i);
|
|
334
|
-
v = e(v.length,
|
|
334
|
+
v = e(v.length, m);
|
|
335
335
|
}
|
|
336
|
-
u("max", v,
|
|
336
|
+
u("max", v, m);
|
|
337
337
|
let [C, w] = l(S.shape, v), E = C;
|
|
338
|
-
|
|
338
|
+
p && (E = o(C, h));
|
|
339
339
|
let D;
|
|
340
340
|
if (x) {
|
|
341
|
-
let e = i.texData.get(S.dataId).values, t =
|
|
341
|
+
let e = i.texData.get(S.dataId).values, t = _(e, c(w), E, d.dtype);
|
|
342
342
|
D = i.makeTensorInfo(E, d.dtype);
|
|
343
343
|
let n = i.texData.get(D.dataId);
|
|
344
344
|
n.values = t;
|
|
@@ -349,7 +349,7 @@ var j = {
|
|
|
349
349
|
kernelName: "Max",
|
|
350
350
|
backendName: "webgl",
|
|
351
351
|
kernelFunc: A
|
|
352
|
-
}, M =
|
|
352
|
+
}, M = m({
|
|
353
353
|
opSnippet: "\nif (a == b) {\n return 1.0;\n};\nreturn a / b;",
|
|
354
354
|
packedOpSnippet: "\n // vec4 one = vec4(equal(a, b));\n // return one + (vec4(1.0) - one) * a / b;\n vec4 result = a / b;\n if(a.x == b.x) {\n result.x = 1.;\n }\n if(a.y == b.y) {\n result.y = 1.;\n }\n if(a.z == b.z) {\n result.z = 1.;\n }\n if(a.w == b.w) {\n result.w = 1.;\n }\n\n return result;\n",
|
|
355
355
|
checkOutOfBounds: !0
|