@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.
Files changed (109) hide show
  1. package/README.md +78 -281
  2. package/dist/{DatasetBuilder-C0iJT29K.js → DatasetBuilder-DU1G1OKX.js} +20 -23
  3. package/dist/{RealDiv-CNsvC4AU.js → RealDiv-CSnvtN2E.js} +20 -20
  4. package/dist/{Reshape-dnm9bO3B.js → Reshape-BlylqwWy.js} +12 -12
  5. package/dist/TeachableLLM.d.ts +10 -15
  6. package/dist/TeachableLLM.js +201 -2
  7. package/dist/api/responses.d.ts +81 -0
  8. package/dist/api/responses.js +169 -0
  9. package/dist/api/training.d.ts +70 -0
  10. package/dist/api/training.js +205 -0
  11. package/dist/data/docx.js +9 -3036
  12. package/dist/data/stream.d.ts +8 -8
  13. package/dist/data/stream.js +1 -1
  14. package/dist/data/textLoader.d.ts +1 -1
  15. package/dist/data/textLoader.js +2 -2
  16. package/dist/data.d.ts +3 -0
  17. package/dist/data.js +12 -0
  18. package/dist/{dist-BqAU9-yi.js → dist-CwK5S7Ls.js} +2168 -2168
  19. package/dist/{gpgpu_math-DBYEAAdI.js → gpgpu_math-20tPK8LM.js} +458 -458
  20. package/dist/{Generator.d.ts → inference/Generator.d.ts} +16 -44
  21. package/dist/inference/Generator.js +271 -0
  22. package/dist/inference/tokenisePrompt.d.ts +4 -0
  23. package/dist/inference/tokenisePrompt.js +13 -0
  24. package/dist/inference/types.d.ts +44 -8
  25. package/dist/inference/utilities.d.ts +9 -0
  26. package/dist/inference/utilities.js +20 -0
  27. package/dist/jszip.min-DKa1Rjyn.js +3033 -0
  28. package/dist/{kernel_funcs_utils-D-mATnGR.js → kernel_funcs_utils-ql8Y8qPn.js} +96 -93
  29. package/dist/layers/MLP.d.ts +1 -1
  30. package/dist/layers/PositionEmbedding.d.ts +2 -1
  31. package/dist/layers/PositionEmbedding.js +1 -1
  32. package/dist/layers/RMSNorm.d.ts +1 -1
  33. package/dist/layers/TiedEmbedding.js +1 -1
  34. package/dist/layers.d.ts +4 -0
  35. package/dist/layers.js +14 -0
  36. package/dist/loader/load.js +58 -2
  37. package/dist/loader/loadHF.d.ts +1 -1
  38. package/dist/loader/loadHF.js +17 -2
  39. package/dist/loader/loadTransformers.js +46 -2
  40. package/dist/loader/newZipLoad.js +25 -2
  41. package/dist/loader/oldZipLoad.d.ts +1 -1
  42. package/dist/loader/oldZipLoad.js +37 -2
  43. package/dist/loader/save.js +75 -2
  44. package/dist/loader/types.d.ts +3 -3
  45. package/dist/main.d.ts +34 -43
  46. package/dist/main.js +12327 -20
  47. package/dist/{matMulGelu-BAIgQaRx.js → matMulGelu-CBoqTZM7.js} +2 -2
  48. package/dist/models/NanoGPTV1.js +95 -2
  49. package/dist/models/NanoGPTV2.js +86 -2
  50. package/dist/models/factory.js +13 -2
  51. package/dist/models/model.js +76 -2
  52. package/dist/models.d.ts +4 -0
  53. package/dist/models.js +14 -0
  54. package/dist/ops/dot16.js +1 -1
  55. package/dist/ops/matMulGelu.js +1 -1
  56. package/dist/ops/webgl/adamAdjust.js +1 -1
  57. package/dist/ops/webgl/fusedSoftmax.js +2 -2
  58. package/dist/ops/webgl/gelu.js +2 -2
  59. package/dist/ops/webgl/log.js +5 -5
  60. package/dist/ops/webgl/matMulGelu.js +1 -1
  61. package/dist/ops/webgl/matMulMul.js +1 -1
  62. package/dist/{stream-BjdpSNqB.js → stream-BpAwcvHz.js} +565 -561
  63. package/dist/{tfjs_backend-CydPRQTc.js → tfjs_backend-h5weiy1O.js} +36 -36
  64. package/dist/tokenise.d.ts +4 -0
  65. package/dist/tokenise.js +15 -0
  66. package/dist/tokeniser/CharTokeniser.js +18 -20
  67. package/dist/tokeniser/bpe.js +18 -22
  68. package/dist/training/BasicTrainer.d.ts +5 -10
  69. package/dist/training/BasicTrainer.js +80 -88
  70. package/dist/training/DatasetBuilder.d.ts +4 -4
  71. package/dist/training/DatasetBuilder.js +1 -1
  72. package/dist/training/PreTrainer.js +1 -1
  73. package/dist/training/SFTTrainer.js +1 -1
  74. package/dist/training/configure.d.ts +3 -0
  75. package/dist/training/configure.js +32 -0
  76. package/dist/training/factory.d.ts +6 -0
  77. package/dist/training/factory.js +8 -0
  78. package/dist/training/prepareData.d.ts +22 -0
  79. package/dist/training/prepareData.js +49 -0
  80. package/dist/training/tasks/TokenStore.d.ts +2 -1
  81. package/dist/training/tasks/TokenStore.js +8 -5
  82. package/dist/training/tasks/tokenStream.d.ts +17 -0
  83. package/dist/training/tasks/tokenStream.js +46 -0
  84. package/dist/training/types.d.ts +14 -1
  85. package/dist/training/validateOptions.d.ts +2 -0
  86. package/dist/training/validateOptions.js +19 -0
  87. package/dist/training/validation.js +4 -2
  88. package/dist/utilities/arrayShape.d.ts +1 -0
  89. package/dist/utilities/arrayShape.js +8 -0
  90. package/dist/utilities/random.d.ts +1 -0
  91. package/dist/utilities/random.js +19 -0
  92. package/dist/utilities/waitForModel.d.ts +1 -1
  93. package/dist/v4-BK7K-jy_.js +30 -0
  94. package/package.json +8 -2
  95. package/dist/Generator.js +0 -2
  96. package/dist/Trainer-DBsyWJ4s.js +0 -228
  97. package/dist/Trainer.d.ts +0 -45
  98. package/dist/Trainer.js +0 -2
  99. package/dist/main-BSaDGH7I.js +0 -13274
  100. package/dist/training/tasks/ConversationTask.d.ts +0 -17
  101. package/dist/training/tasks/ConversationTask.js +0 -29
  102. package/dist/training/tasks/PretrainingTask.d.ts +0 -17
  103. package/dist/training/tasks/PretrainingTask.js +0 -42
  104. package/dist/training/tasks/StartSentenceTask.d.ts +0 -18
  105. package/dist/training/tasks/StartSentenceTask.js +0 -45
  106. package/dist/training/tasks/Task.d.ts +0 -29
  107. package/dist/training/tasks/Task.js +0 -50
  108. package/dist/training/tasks/splitter.d.ts +0 -5
  109. 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 GPT language models built on TensorFlow.js, developed as part of the Finnish Generation AI research project. This library enables training, fine-tuning, and inference of transformer-based language models entirely in the browser with support for explainable AI (XAI) features. It is intended to be used as an educational tool for learning about the model training process since it targets mostly tiny models. In principle it could be adapted to load other pre-trained models from Hugging Face.
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 version available here: https://lm.gen-ai.fi
5
+ Live demo: https://lm.gen-ai.fi
6
6
 
7
- ## Overview
7
+ **Design goals**
8
8
 
9
- GenAI NanoGPT is inspired by [Andrej Karpathy's NanoGPT](https://github.com/karpathy/nanoGPT) but reimagined for the browser using TensorFlow.js. It provides a complete pipeline for:
10
-
11
- - **Training** language models from scratch in the browser
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
- ## Quick Start
19
+ ## Main concepts
35
20
 
36
- ### Creating and Training a Model
37
-
38
- ```javascript
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
- // Select the best available backend
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
- // Create a new model
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
- // Training data
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, waitForModel } from '@genai-fi/nanogpt';
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
- ## Event Handlers and Real-time Updates
102
-
103
- ### Monitoring Training Progress
104
-
105
- Track training metrics in real-time with event handlers:
106
-
107
- ```javascript
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
- await model.train(trainingText, options);
43
+ // Switch backend if needed
44
+ await TeachableLLM.selectBackend('webgpu');
124
45
  ```
125
46
 
126
- ### Real-time Token Generation
127
-
128
- Stream generated tokens as they're produced:
47
+ Training the tokeniser (when using streamed conversational data)
129
48
 
130
49
  ```javascript
131
- const generator = model.generator();
50
+ import { data, tokenise } from '@genai-fi/nanogpt';
132
51
 
133
- // Listen for generated tokens
134
- generator.on('tokens', (tokens) => {
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
- // Update UI incrementally
140
- appendToOutput(text);
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
- ## Training on Small Devices
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
- // Smaller model configuration for mobile devices
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
- // Training options optimized for limited memory
182
- await mobileModel.train(trainingText, {
183
- batchSize: 8, // Smaller batch size
184
- learningRate: 3e-4,
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
- 1. **Start Small**: Use smaller models (4 layers) and shorter context windows (128 tokens)
196
- 2. **Reduce Batch Size**: Use batch sizes of 8-16 depending on available memory
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
- ### Attention Visualization
74
+ Generate text (responses API)
203
75
 
204
76
  ```javascript
205
- const generator = model.generator();
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
- const text = await generator.generate('Prompt', {
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
- const probabilities = generator.getProbabilitiesData();
217
- // Shape: [num_tokens][seq_len][vocab_size]
92
+ // Manage responses
93
+ // model.responses.cancel(id)
94
+ // model.responses.hook(id)
95
+ // model.responses.resume(id)
218
96
  ```
219
97
 
220
- ### Streaming Generation
98
+ Tokenisers and token stores
221
99
 
222
100
  ```javascript
223
- const generator = model.generator();
101
+ import { tokenise } from '@genai-fi/nanogpt';
224
102
 
225
- generator.on('tokens', (tokens) => {
226
- // Update UI with new tokens in real-time
227
- updateDisplay(tokens);
228
- });
103
+ // Character and BPE tokenisers are available
104
+ const { CharTokeniser, BPETokeniser, TokenStore, createTokenStore } = tokenise;
229
105
 
230
- generator.on('start', () => console.log('Generation started'));
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
- ### Memory Management
109
+ Data helpers
239
110
 
240
111
  ```javascript
241
- // Enable profiling
242
- model.enableProfiler = true;
112
+ import { data } from '@genai-fi/nanogpt';
243
113
 
244
- // After training/generation
245
- const profiler = model.getProfiler();
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
- // Clean up
251
- model.dispose();
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
- ### Setup
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
- ### Building
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 lint # Check code style
134
+ npm run build
135
+ npm run dev
136
+ npm test
137
+ npm run test:gl
336
138
  ```
337
139
 
338
- ## Performance Tips
140
+ ## Examples and demos
339
141
 
340
- 1. **Use WebGPU**: Provides the best performance for training and inference
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
- - Inspired by [Andrej Karpathy's NanoGPT](https://github.com/karpathy/nanoGPT)
350
- - Built with [TensorFlow.js](https://www.tensorflow.org/js)
351
- - Developed as part of the Finnish [Generation AI research project](https://generation-ai-stn.fi)
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 || S(e) || Array.isArray(e) || typeof e == "object" && e instanceof f || l(e);
61
+ return e == null || ee(e) || Array.isArray(e) || typeof e == "object" && e instanceof f || l(e);
62
62
  }
63
- function S(e) {
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 C(e) {
69
- return h(e, ee);
68
+ function S(e) {
69
+ return h(e, C);
70
70
  }
71
- function ee(e) {
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: C(e),
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 Uint32Array(e.getShardCount()),
687
- shuffledIndexes: new Uint32Array(o),
688
- lastShardIndexes: new Uint32Array(Math.ceil(e.getShardLength(e.getShardCount() - 1) / this.blockSize)),
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], n = t * this.blockSize, a = c.currentShard, o = c.currentMask, s = $(c, e, r);
707
- if (!a) break;
708
- if (n + this.blockSize + 1 > a.length) {
709
- console.warn("Index out of bounds for current shard, moving to next shard", t, n, a.length);
710
- continue;
711
- }
712
- let l = new Int32Array(a.slice(n, n + this.blockSize)), u = new Int32Array(a.slice(n + 1, n + this.blockSize + 1));
713
- if (o) {
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: l,
720
- ys: u
721
- }, await s;
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 { o as f } from "./gpgpu_math-DBYEAAdI.js";
4
- import { J as p, b as m } from "./shared-CfTzpULd.js";
5
- import { S as h, n as g } from "./kernel_funcs_utils-D-mATnGR.js";
6
- import { t as _ } from "./Reshape-dnm9bO3B.js";
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 = h("rc", this.rank), a = Array(this.rank);
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), h = m, g = t(h, p), v = g != null, y = n;
282
- v && (y = T(n, g, d), h = e(h.length, p)), u("sum", h, p);
283
- let [b, S] = l(y.shape, h), C = b;
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: h } = a, g = d.shape.length, _ = s(f, d.shape), v = _, y = t(v, g), b = y != null, x = i.shouldExecuteOnCPU([d]), S = d;
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(g);
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 = p(e, d.shape, d.dtype, y, t);
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, g);
334
+ v = e(v.length, m);
335
335
  }
336
- u("max", v, g);
336
+ u("max", v, m);
337
337
  let [C, w] = l(S.shape, v), E = C;
338
- h && (E = o(C, _));
338
+ p && (E = o(C, h));
339
339
  let D;
340
340
  if (x) {
341
- let e = i.texData.get(S.dataId).values, t = m(e, c(w), E, d.dtype);
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 = g({
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