edgeinfer 0.0.0-stage → 2.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/dist/index.js ADDED
@@ -0,0 +1,1827 @@
1
+ "use strict";
2
+ var __create = Object.create;
3
+ var __defProp = Object.defineProperty;
4
+ var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
5
+ var __getOwnPropNames = Object.getOwnPropertyNames;
6
+ var __getProtoOf = Object.getPrototypeOf;
7
+ var __hasOwnProp = Object.prototype.hasOwnProperty;
8
+ var __export = (target, all) => {
9
+ for (var name in all)
10
+ __defProp(target, name, { get: all[name], enumerable: true });
11
+ };
12
+ var __copyProps = (to, from, except, desc) => {
13
+ if (from && typeof from === "object" || typeof from === "function") {
14
+ for (let key of __getOwnPropNames(from))
15
+ if (!__hasOwnProp.call(to, key) && key !== except)
16
+ __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
17
+ }
18
+ return to;
19
+ };
20
+ var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__getProtoOf(mod)) : {}, __copyProps(
21
+ // If the importer is in node compatibility mode or this is not an ESM
22
+ // file that has been converted to a CommonJS file using a Babel-
23
+ // compatible transform (i.e. "__esModule" has not been set), then set
24
+ // "default" to the CommonJS "module.exports" for node compatibility.
25
+ isNodeMode || !mod || !mod.__esModule ? __defProp(target, "default", { value: mod, enumerable: true }) : target,
26
+ mod
27
+ ));
28
+ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
29
+
30
+ // src/index.ts
31
+ var index_exports = {};
32
+ __export(index_exports, {
33
+ EdgeInfer: () => EdgeInfer,
34
+ EventEmitter: () => EventEmitter,
35
+ ImageProcessor: () => ImageProcessor,
36
+ IndexedDbVectorStore: () => IndexedDbVectorStore,
37
+ MatryoshkaIndex: () => MatryoshkaIndex,
38
+ MemoryVectorStore: () => MemoryVectorStore,
39
+ MockLanguageModel: () => MockLanguageModel,
40
+ ModelCache: () => ModelCache,
41
+ OnnxCausalLM: () => OnnxCausalLM,
42
+ RuntimeManager: () => RuntimeManager,
43
+ SpeculativeDecoder: () => SpeculativeDecoder,
44
+ Tokenizer: () => Tokenizer,
45
+ createMatryoshkaEmbedder: () => createMatryoshkaEmbedder,
46
+ dequantizeInt8: () => dequantizeInt8,
47
+ embedMatryoshka: () => embedMatryoshka,
48
+ l2Normalize: () => l2Normalize,
49
+ probability: () => probability,
50
+ quantizeInt8: () => quantizeInt8,
51
+ residualDistribution: () => residualDistribution,
52
+ sample: () => sample,
53
+ softmax: () => softmax,
54
+ truncateAndNormalize: () => truncateAndNormalize
55
+ });
56
+ module.exports = __toCommonJS(index_exports);
57
+
58
+ // src/events.ts
59
+ var EventEmitter = class {
60
+ constructor() {
61
+ this.listeners = {};
62
+ }
63
+ on(event, callback) {
64
+ if (!this.listeners[event]) this.listeners[event] = [];
65
+ this.listeners[event].push(callback);
66
+ }
67
+ emit(event, ...args) {
68
+ if (this.listeners[event]) {
69
+ this.listeners[event].forEach((cb) => cb(...args));
70
+ }
71
+ }
72
+ };
73
+
74
+ // src/download.ts
75
+ function checkDownloadAbort(signal) {
76
+ if (signal?.aborted) throw signal.reason ?? new DOMException("Download aborted", "AbortError");
77
+ }
78
+ async function withDownloadDeadline(operation, options = {}) {
79
+ const timeout = options.downloadTimeoutMs ?? 12e4;
80
+ if (!Number.isFinite(timeout) || timeout <= 0 || timeout > 2147483647) {
81
+ throw new RangeError("downloadTimeoutMs must be between 0 and 2147483647 milliseconds");
82
+ }
83
+ const controller = new AbortController();
84
+ const abort = () => controller.abort(options.signal?.reason);
85
+ if (options.signal?.aborted) abort();
86
+ else options.signal?.addEventListener("abort", abort, { once: true });
87
+ const timer = setTimeout(() => controller.abort(new DOMException("Download deadline exceeded", "TimeoutError")), timeout);
88
+ let onAbort = () => {
89
+ };
90
+ try {
91
+ checkDownloadAbort(controller.signal);
92
+ const cancelled = new Promise((_, reject) => {
93
+ onAbort = () => reject(controller.signal.reason ?? new DOMException("Download aborted", "AbortError"));
94
+ controller.signal.addEventListener("abort", onAbort, { once: true });
95
+ });
96
+ return await Promise.race([operation(controller.signal), cancelled]);
97
+ } finally {
98
+ clearTimeout(timer);
99
+ controller.signal.removeEventListener("abort", onAbort);
100
+ options.signal?.removeEventListener("abort", abort);
101
+ }
102
+ }
103
+
104
+ // src/model-cache.ts
105
+ var ModelCache = class extends EventEmitter {
106
+ constructor(cacheName = "edgeinfer-models") {
107
+ super();
108
+ this.cacheName = cacheName;
109
+ }
110
+ async getModel(url, forceDownload = false, options = {}) {
111
+ return withDownloadDeadline(async (signal) => {
112
+ const cache = typeof caches === "undefined" ? void 0 : await caches.open(this.cacheName);
113
+ checkDownloadAbort(signal);
114
+ const cached = !forceDownload && cache ? await cache.match(url) : void 0;
115
+ checkDownloadAbort(signal);
116
+ if (cached) {
117
+ const buffer2 = await cached.arrayBuffer();
118
+ checkDownloadAbort(signal);
119
+ return buffer2;
120
+ }
121
+ this.emit("downloadStart", url);
122
+ checkDownloadAbort(signal);
123
+ const response = await fetch(url, { signal });
124
+ if (!response.ok) {
125
+ await response.body?.cancel();
126
+ throw new Error(`EdgeInfer: failed to fetch model from ${url} - HTTP ${response.status} ${response.statusText}`);
127
+ }
128
+ if (response.status === 206 || response.headers.has("content-range")) {
129
+ await response.body?.cancel();
130
+ throw new Error("EdgeInfer: partial model responses are not supported");
131
+ }
132
+ const buffer = await response.arrayBuffer();
133
+ checkDownloadAbort(signal);
134
+ if (cache) await cache.put(url, new Response(buffer, { headers: response.headers }));
135
+ checkDownloadAbort(signal);
136
+ this.emit("downloadComplete", url);
137
+ return buffer;
138
+ }, options);
139
+ }
140
+ };
141
+
142
+ // src/runtime-manager.ts
143
+ var ort = __toESM(require("onnxruntime-web"));
144
+ var ADAPTER_TIMEOUT_MS = 5e3;
145
+ var _RuntimeManager = class _RuntimeManager {
146
+ /**
147
+ * Detect GPU capabilities and recommend the best execution configuration.
148
+ */
149
+ static async detectCapabilities() {
150
+ if (this.cachedCapabilities) return this.cachedCapabilities;
151
+ const capabilities = {
152
+ provider: "wasm",
153
+ hasWebGPU: false,
154
+ hasWebNN: false,
155
+ estimatedVRAM: 0,
156
+ adapterLimits: null,
157
+ vramProbe: null,
158
+ recommendedQuantization: "int8"
159
+ };
160
+ if (typeof navigator !== "undefined" && "gpu" in navigator) {
161
+ try {
162
+ const gpu = navigator.gpu;
163
+ const adapter = await _RuntimeManager.withTimeout(
164
+ Promise.resolve(gpu.requestAdapter()),
165
+ ADAPTER_TIMEOUT_MS
166
+ );
167
+ if (adapter) {
168
+ capabilities.hasWebGPU = true;
169
+ capabilities.provider = "webgpu";
170
+ const maxBufferSize = adapter.limits?.maxBufferSize || 0;
171
+ const maxStorageSize = adapter.limits?.maxStorageBufferBindingSize || 0;
172
+ capabilities.adapterLimits = {
173
+ maxBufferSize,
174
+ maxStorageBufferBindingSize: maxStorageSize
175
+ };
176
+ const probe = await _RuntimeManager.probeVRAM(adapter);
177
+ capabilities.vramProbe = probe;
178
+ capabilities.estimatedVRAM = probe.available ? probe.maxAllocationBytes : Math.max(maxBufferSize, maxStorageSize);
179
+ const vramGB = capabilities.estimatedVRAM / (1024 * 1024 * 1024);
180
+ if (vramGB >= 4) {
181
+ capabilities.recommendedQuantization = "fp16";
182
+ } else if (vramGB >= 1) {
183
+ capabilities.recommendedQuantization = "int8";
184
+ } else {
185
+ capabilities.recommendedQuantization = "int4";
186
+ }
187
+ }
188
+ } catch {
189
+ }
190
+ }
191
+ if (typeof navigator !== "undefined" && "ml" in navigator) {
192
+ try {
193
+ const ml = navigator.ml;
194
+ if (ml) {
195
+ capabilities.hasWebNN = true;
196
+ if (!capabilities.hasWebGPU) {
197
+ capabilities.provider = "webnn";
198
+ }
199
+ }
200
+ } catch {
201
+ }
202
+ }
203
+ this.cachedCapabilities = capabilities;
204
+ return capabilities;
205
+ }
206
+ /**
207
+ * Get the best available execution provider with smart fallback.
208
+ * Priority: WebGPU → WebNN → WASM
209
+ */
210
+ static async getBestExecutionProvider() {
211
+ const caps = await this.detectCapabilities();
212
+ return caps.provider;
213
+ }
214
+ /**
215
+ * Create an ONNX inference session with automatic provider selection.
216
+ * Falls back gracefully if the preferred provider fails.
217
+ *
218
+ * Returns BOTH the session and the provider that actually served it. The
219
+ * detected capability (`detectCapabilities().provider`) is only a preference:
220
+ * a provider can be detected as available and still fail to create a session
221
+ * (for example onnxruntime-web only registers the WebGPU execution provider
222
+ * in some builds/versions). Callers must report the returned `provider`, not
223
+ * the detected one, or `EdgeInfer.executionProvider` will lie.
224
+ */
225
+ static async createSession(modelBuffer, providers) {
226
+ const requestedProviders = providers || [await this.getBestExecutionProvider()];
227
+ for (const provider of requestedProviders) {
228
+ try {
229
+ const session2 = await ort.InferenceSession.create(
230
+ modelBuffer,
231
+ {
232
+ executionProviders: [provider],
233
+ graphOptimizationLevel: "all"
234
+ }
235
+ );
236
+ return { session: session2, provider };
237
+ } catch (err) {
238
+ console.warn(`EdgeInfer: Failed to create session with provider "${provider}", trying next...`, err);
239
+ }
240
+ }
241
+ console.warn("EdgeInfer: All preferred providers failed, falling back to WASM");
242
+ const session = await ort.InferenceSession.create(modelBuffer, {
243
+ executionProviders: ["wasm"],
244
+ graphOptimizationLevel: "all"
245
+ });
246
+ return { session, provider: "wasm" };
247
+ }
248
+ /**
249
+ * Empirically probe the largest single storage buffer the adapter's device
250
+ * will allocate. Allocates and destroys one buffer at a time; never relies on
251
+ * `maxBufferSize` as if it were physical VRAM.
252
+ */
253
+ static async probeVRAM(adapter, options = {}) {
254
+ const limits = {
255
+ maxBufferSize: Number(adapter?.limits?.maxBufferSize) || 0,
256
+ maxStorageBufferBindingSize: Number(adapter?.limits?.maxStorageBufferBindingSize) || 0
257
+ };
258
+ const unavailable = {
259
+ available: false,
260
+ maxAllocationBytes: 0,
261
+ testedAllocations: [],
262
+ failedAllocationBytes: null,
263
+ adapterLimits: limits,
264
+ deviceCreated: false,
265
+ probe: "unavailable"
266
+ };
267
+ if (!adapter || typeof adapter.requestDevice !== "function") return unavailable;
268
+ let device;
269
+ try {
270
+ device = await adapter.requestDevice();
271
+ } catch {
272
+ return unavailable;
273
+ }
274
+ if (!device || typeof device.createBuffer !== "function") {
275
+ return { ...unavailable, deviceCreated: Boolean(device) };
276
+ }
277
+ const MIB = 1024 * 1024;
278
+ const start = Math.max(1, options.startBytes ?? 64 * MIB);
279
+ const configuredMax = Math.max(start, options.maxBytes ?? 4 * 1024 * MIB);
280
+ const max = limits.maxBufferSize > 0 ? Math.min(configuredMax, limits.maxBufferSize) : configuredMax;
281
+ const step = Math.max(1, options.stepBytes ?? 32 * MIB);
282
+ const tested = [];
283
+ const STORAGE = 128;
284
+ const COPY_DST = 8;
285
+ const attempt = (size2) => {
286
+ tested.push(size2);
287
+ try {
288
+ const buffer = device.createBuffer({ size: size2, usage: STORAGE | COPY_DST });
289
+ buffer?.destroy?.();
290
+ return true;
291
+ } catch {
292
+ return false;
293
+ }
294
+ };
295
+ let low = 0;
296
+ let high = 0;
297
+ let size = Math.min(start, max);
298
+ while (size >= 1) {
299
+ if (attempt(size)) {
300
+ low = size;
301
+ break;
302
+ }
303
+ high = size;
304
+ if (size === 1) break;
305
+ size = Math.max(1, Math.floor(size / 2));
306
+ }
307
+ if (low === 0) {
308
+ device.destroy?.();
309
+ return {
310
+ ...unavailable,
311
+ testedAllocations: tested,
312
+ failedAllocationBytes: high || null,
313
+ deviceCreated: true,
314
+ probe: "limits"
315
+ };
316
+ }
317
+ let next = Math.min(max, low * 2);
318
+ while (next > low && tested.length < 16) {
319
+ if (attempt(next)) {
320
+ low = next;
321
+ next = Math.min(max, low * 2);
322
+ if (low === max) break;
323
+ } else {
324
+ high = next;
325
+ break;
326
+ }
327
+ }
328
+ if (high === 0) high = Math.min(max, low * 2);
329
+ while (high - low > step && tested.length < 24) {
330
+ const mid = Math.floor((low + high) / 2);
331
+ if (mid <= low || mid >= high) break;
332
+ if (attempt(mid)) low = mid;
333
+ else high = mid;
334
+ }
335
+ device.destroy?.();
336
+ return {
337
+ available: true,
338
+ maxAllocationBytes: low,
339
+ testedAllocations: tested,
340
+ failedAllocationBytes: high > low ? high : null,
341
+ adapterLimits: limits,
342
+ deviceCreated: true,
343
+ probe: "empirical"
344
+ };
345
+ }
346
+ /**
347
+ * Resolve `promise`, or resolve to `null` if it has not settled within `ms`.
348
+ * Used to stop a hung GPU driver from blocking model loading forever.
349
+ */
350
+ static withTimeout(promise, ms) {
351
+ let timer;
352
+ return Promise.race([
353
+ promise.then((v) => {
354
+ clearTimeout(timer);
355
+ return v;
356
+ }),
357
+ new Promise((resolve) => {
358
+ timer = setTimeout(() => {
359
+ console.warn(`EdgeInfer: navigator.gpu.requestAdapter() did not settle within ${ms}ms; falling back to WASM.`);
360
+ resolve(null);
361
+ }, ms);
362
+ if (typeof timer?.unref === "function") timer.unref();
363
+ })
364
+ ]);
365
+ }
366
+ /**
367
+ * Estimate if a model of given size (bytes) can fit in available VRAM.
368
+ */
369
+ static async canFitModel(modelSizeBytes) {
370
+ const caps = await this.detectCapabilities();
371
+ if (!caps.hasWebGPU) return true;
372
+ return caps.estimatedVRAM >= modelSizeBytes * 1.5;
373
+ }
374
+ /**
375
+ * Reset cached capabilities (useful for testing or re-detection).
376
+ */
377
+ static resetCache() {
378
+ this.cachedCapabilities = null;
379
+ }
380
+ };
381
+ _RuntimeManager.cachedCapabilities = null;
382
+ var RuntimeManager = _RuntimeManager;
383
+
384
+ // src/tokenizer.ts
385
+ var Tokenizer = class _Tokenizer {
386
+ constructor(config) {
387
+ this.vocab = new Map(Object.entries(config.vocab));
388
+ this.inverseVocab = /* @__PURE__ */ new Map();
389
+ for (const [token, id] of this.vocab) {
390
+ this.inverseVocab.set(id, token);
391
+ }
392
+ this.merges = (config.merges || []).map((m) => {
393
+ if (Array.isArray(m)) {
394
+ return [m[0], m[1]];
395
+ }
396
+ const parts = String(m).split(" ");
397
+ return [parts[0], parts[1]];
398
+ });
399
+ this.specialTokens = new Map(Object.entries(config.specialTokens || {}));
400
+ this.maxLength = config.maxLength || 512;
401
+ this.padTokenId = config.padTokenId ?? 0;
402
+ this.unkTokenId = config.unkTokenId ?? 100;
403
+ this.clsTokenId = config.clsTokenId ?? 101;
404
+ this.sepTokenId = config.sepTokenId ?? 102;
405
+ this.bpeCache = /* @__PURE__ */ new Map();
406
+ this.usesBPE = this.merges.length > 0;
407
+ this.bpeSpaceMarker = config.bpeSpaceMarker ?? "";
408
+ this.mergeRanks = /* @__PURE__ */ new Map();
409
+ this.merges.forEach(([first, second], rank) => {
410
+ this.mergeRanks.set(`${first} ${second}`, rank);
411
+ });
412
+ this.lowercase = config.lowercase ?? !this.usesBPE;
413
+ }
414
+ /**
415
+ * Load tokenizer configuration from a JSON URL (tokenizer.json or vocab.json).
416
+ */
417
+ static async fromUrl(url, options = {}) {
418
+ const config = await withDownloadDeadline(async (signal) => {
419
+ const response = await fetch(url, { signal });
420
+ if (!response.ok) {
421
+ await response.body?.cancel();
422
+ throw new Error(`Failed to load tokenizer from ${url}: ${response.status}`);
423
+ }
424
+ return response.json();
425
+ }, options);
426
+ if (config.model && config.model.vocab) {
427
+ const specialTokens = {};
428
+ if (config.added_tokens) {
429
+ for (const t of config.added_tokens) {
430
+ specialTokens[t.content] = t.id;
431
+ }
432
+ }
433
+ return new _Tokenizer({
434
+ vocab: config.model.vocab,
435
+ lowercase: config.normalizer === void 0 ? void 0 : _Tokenizer.normalizerLowercase(config.normalizer),
436
+ merges: config.model.merges || [],
437
+ specialTokens,
438
+ maxLength: config.truncation?.max_length || 512,
439
+ padTokenId: specialTokens["[PAD]"] ?? 0,
440
+ unkTokenId: specialTokens["[UNK]"] ?? 100,
441
+ clsTokenId: specialTokens["[CLS]"] ?? 101,
442
+ sepTokenId: specialTokens["[SEP]"] ?? 102
443
+ });
444
+ }
445
+ return new _Tokenizer({ vocab: config });
446
+ }
447
+ static normalizerLowercase(normalizer) {
448
+ if (normalizer === void 0) return true;
449
+ if (normalizer === null) return false;
450
+ if (typeof normalizer !== "object") throw new Error("Unsupported tokenizer normalizer");
451
+ const value = normalizer;
452
+ if (value.type === "Lowercase") return true;
453
+ if (value.type === "BertNormalizer" && typeof value.lowercase === "boolean" && value.strip_accents === false && value.clean_text === false && value.handle_chinese_chars === false) {
454
+ return value.lowercase;
455
+ }
456
+ throw new Error("Unsupported tokenizer normalizer; only identity, Lowercase, or casing-only BertNormalizer is supported");
457
+ }
458
+ /**
459
+ * Encode text into token IDs with attention mask.
460
+ * Handles WordPiece tokenization for BERT-like models.
461
+ */
462
+ encode(text, options) {
463
+ const addSpecial = options?.addSpecialTokens !== false;
464
+ const maxLen = options?.maxLength || this.maxLength;
465
+ const padding = options?.padding ?? false;
466
+ const words = this.preTokenize(text);
467
+ const tokenIds = [];
468
+ if (addSpecial) {
469
+ tokenIds.push(this.clsTokenId);
470
+ }
471
+ for (const word of words) {
472
+ const subTokens = this.tokenizeWord(word);
473
+ for (const subToken of subTokens) {
474
+ if (tokenIds.length >= maxLen - (addSpecial ? 1 : 0)) break;
475
+ tokenIds.push(subToken);
476
+ }
477
+ }
478
+ if (addSpecial) {
479
+ tokenIds.push(this.sepTokenId);
480
+ }
481
+ const truncated = tokenIds.slice(0, maxLen);
482
+ if (truncated.length === 0) {
483
+ const fallback = new Int32Array([this.padTokenId]);
484
+ return {
485
+ inputIds: fallback,
486
+ attentionMask: new Int32Array([0]),
487
+ inputIdsInt64: _Tokenizer.toInt64(fallback),
488
+ attentionMaskInt64: new BigInt64Array([0n])
489
+ };
490
+ }
491
+ const attentionMask = new Int32Array(padding ? maxLen : truncated.length);
492
+ attentionMask.fill(1, 0, truncated.length);
493
+ const inputIds = new Int32Array(padding ? maxLen : truncated.length);
494
+ inputIds.set(truncated);
495
+ if (padding) {
496
+ inputIds.fill(this.padTokenId, truncated.length);
497
+ }
498
+ return {
499
+ inputIds,
500
+ attentionMask,
501
+ inputIdsInt64: _Tokenizer.toInt64(inputIds),
502
+ attentionMaskInt64: _Tokenizer.toInt64(attentionMask)
503
+ };
504
+ }
505
+ /**
506
+ * Decode token IDs back to text.
507
+ */
508
+ decode(tokens) {
509
+ const tokenArray = tokens instanceof Int32Array ? Array.from(tokens) : tokens;
510
+ const words = [];
511
+ for (const id of tokenArray) {
512
+ if (id === this.padTokenId || id === this.clsTokenId || id === this.sepTokenId) continue;
513
+ const token = this.inverseVocab.get(id);
514
+ if (!token) {
515
+ words.push("[UNK]");
516
+ } else if (token.startsWith("##")) {
517
+ if (words.length > 0) {
518
+ words[words.length - 1] += token.slice(2);
519
+ } else {
520
+ words.push(token.slice(2));
521
+ }
522
+ } else if (token.startsWith("\u2581") || token.startsWith("\u0120")) {
523
+ words.push(token.slice(1));
524
+ } else {
525
+ words.push(token);
526
+ }
527
+ }
528
+ return words.join(" ");
529
+ }
530
+ /**
531
+ * Get vocabulary size.
532
+ */
533
+ get vocabSize() {
534
+ return this.vocab.size;
535
+ }
536
+ /**
537
+ * Resolve a special token (e.g. `</s>`, `<|endoftext|>`) to its ID, whether it
538
+ * was declared in `specialTokens` or is present in the main vocabulary.
539
+ */
540
+ getSpecialTokenId(token) {
541
+ return this.specialTokens.get(token) ?? this.vocab.get(token);
542
+ }
543
+ /**
544
+ * Pre-tokenize text: lowercase, split on whitespace and punctuation boundaries.
545
+ */
546
+ preTokenize(text) {
547
+ const normalized = (this.lowercase ? text.toLowerCase() : text).trim();
548
+ if (!normalized) return [];
549
+ const tokens = [];
550
+ let current = "";
551
+ for (const char of normalized) {
552
+ if (/\s/.test(char)) {
553
+ if (current) tokens.push(current);
554
+ current = "";
555
+ } else if (/[.,!?;:'"()\[\]{}\-\/\\@#$%^&*~`|<>]/.test(char) || /\p{P}/u.test(char)) {
556
+ if (current) tokens.push(current);
557
+ tokens.push(char);
558
+ current = "";
559
+ } else {
560
+ current += char;
561
+ }
562
+ }
563
+ if (current) tokens.push(current);
564
+ return tokens;
565
+ }
566
+ /**
567
+ * Tokenize a single word using the tokenizer's model type: BPE when a merge
568
+ * table is present, WordPiece otherwise.
569
+ */
570
+ tokenizeWord(word) {
571
+ if (this.usesBPE) return this.tokenizeWordBPE(word);
572
+ if (this.vocab.has(word)) {
573
+ return [this.vocab.get(word)];
574
+ }
575
+ const tokens = [];
576
+ let start = 0;
577
+ while (start < word.length) {
578
+ let end = word.length;
579
+ let found = false;
580
+ while (start < end) {
581
+ let substr = word.slice(start, end);
582
+ if (start > 0) {
583
+ substr = "##" + substr;
584
+ }
585
+ if (this.vocab.has(substr)) {
586
+ tokens.push(this.vocab.get(substr));
587
+ found = true;
588
+ break;
589
+ }
590
+ end--;
591
+ }
592
+ if (!found) {
593
+ tokens.push(this.unkTokenId);
594
+ start++;
595
+ } else {
596
+ start = end;
597
+ }
598
+ }
599
+ return tokens;
600
+ }
601
+ /**
602
+ * Apply the merge table to a word, then map the resulting pieces to IDs.
603
+ * Unknown pieces fall back to the unknown token rather than throwing, which
604
+ * matches ONNX Runtime's text-tokenizer behaviour.
605
+ */
606
+ tokenizeWordBPE(word) {
607
+ const whole = this.vocab.get(
608
+ this.bpeSpaceMarker ? this.bpeSpaceMarker + word : word
609
+ );
610
+ if (whole !== void 0) return [whole];
611
+ const pieces = this.applyBPE(word);
612
+ const ids = [];
613
+ for (let i = 0; i < pieces.length; i++) {
614
+ const piece = i === 0 && this.bpeSpaceMarker ? this.bpeSpaceMarker + pieces[i] : pieces[i];
615
+ const id = this.vocab.get(piece);
616
+ if (id !== void 0) {
617
+ ids.push(id);
618
+ continue;
619
+ }
620
+ ids.push(this.unkTokenId);
621
+ }
622
+ return ids;
623
+ }
624
+ /**
625
+ * Apply BPE merges to a word (for BPE-style tokenizers like GPT-2/Llama).
626
+ * Merge rules are applied in ascending rank order until no adjacent pair is
627
+ * present in the merge table; caching is keyed on the original word.
628
+ */
629
+ applyBPE(word) {
630
+ const cached = this.bpeCache.get(word);
631
+ if (cached) return cached;
632
+ let tokens = word.split("");
633
+ if (tokens.length <= 1 || this.mergeRanks.size === 0) {
634
+ this.bpeCache.set(word, tokens);
635
+ return tokens;
636
+ }
637
+ while (tokens.length > 1) {
638
+ let bestRank = Number.POSITIVE_INFINITY;
639
+ for (let i = 0; i < tokens.length - 1; i++) {
640
+ const rank = this.mergeRanks.get(`${tokens[i]} ${tokens[i + 1]}`);
641
+ if (rank !== void 0 && rank < bestRank) bestRank = rank;
642
+ }
643
+ if (bestRank === Number.POSITIVE_INFINITY) break;
644
+ const merged = [];
645
+ for (let i = 0; i < tokens.length; i++) {
646
+ if (i < tokens.length - 1 && this.mergeRanks.get(`${tokens[i]} ${tokens[i + 1]}`) === bestRank) {
647
+ merged.push(tokens[i] + tokens[i + 1]);
648
+ i++;
649
+ } else {
650
+ merged.push(tokens[i]);
651
+ }
652
+ }
653
+ if (merged.length === tokens.length) break;
654
+ tokens = merged;
655
+ }
656
+ this.bpeCache.set(word, tokens);
657
+ return tokens;
658
+ }
659
+ static toInt64(values) {
660
+ const out = new BigInt64Array(values.length);
661
+ for (let i = 0; i < values.length; i++) out[i] = BigInt(values[i]);
662
+ return out;
663
+ }
664
+ };
665
+
666
+ // src/image-processor.ts
667
+ var ImageProcessor = class {
668
+ static imageDataToFloat32Array(imageData, options = {}) {
669
+ const { normalize = true, layout = "NCHW" } = options;
670
+ const { width, height, data } = imageData;
671
+ const size = width * height;
672
+ const float32Data = new Float32Array(size * 3);
673
+ for (let i = 0; i < size; i++) {
674
+ const r = data[i * 4];
675
+ const g = data[i * 4 + 1];
676
+ const b = data[i * 4 + 2];
677
+ const nr = normalize ? r / 255 : r;
678
+ const ng = normalize ? g / 255 : g;
679
+ const nb = normalize ? b / 255 : b;
680
+ if (layout === "NCHW") {
681
+ float32Data[i] = nr;
682
+ float32Data[size + i] = ng;
683
+ float32Data[2 * size + i] = nb;
684
+ } else {
685
+ float32Data[i * 3] = nr;
686
+ float32Data[i * 3 + 1] = ng;
687
+ float32Data[i * 3 + 2] = nb;
688
+ }
689
+ }
690
+ return float32Data;
691
+ }
692
+ };
693
+
694
+ // src/edgeinfer.ts
695
+ var ort2 = __toESM(require("onnxruntime-web"));
696
+ function throwIfAborted(signal) {
697
+ if (signal?.aborted) {
698
+ throw signal.reason ?? new DOMException("Generation aborted", "AbortError");
699
+ }
700
+ }
701
+ function normalizeChatMessages(messages) {
702
+ if (!Array.isArray(messages) || messages.length === 0) {
703
+ throw new Error("chat requires at least one message");
704
+ }
705
+ return messages.map((message, index) => {
706
+ if (!message || typeof message !== "object") {
707
+ throw new Error(`chat message ${index} must be an object`);
708
+ }
709
+ if (message.role !== "system" && message.role !== "user" && message.role !== "assistant") {
710
+ throw new Error(`chat message ${index} has an invalid role`);
711
+ }
712
+ if (typeof message.content !== "string") {
713
+ throw new Error(`chat message ${index} content must be a string`);
714
+ }
715
+ return { role: message.role, content: message.content };
716
+ });
717
+ }
718
+ function defaultChatTemplate(messages) {
719
+ return `${messages.map((message) => `<|${message.role}|>${message.content}`).join("\n")}
720
+ <|assistant|>`;
721
+ }
722
+ var EdgeInfer = class _EdgeInfer {
723
+ constructor(session, modelSize, provider, int64Inputs, tokenizer) {
724
+ this.tokenizer = null;
725
+ this._provider = "wasm";
726
+ this.session = session;
727
+ this._inputNames = session.inputNames;
728
+ this._outputNames = session.outputNames;
729
+ this._modelSize = modelSize;
730
+ this._provider = provider;
731
+ this._int64Inputs = int64Inputs;
732
+ this.tokenizer = tokenizer || null;
733
+ }
734
+ /**
735
+ * Load a model from URL with automatic caching and optimal provider selection.
736
+ */
737
+ static async load(modelUrl, options) {
738
+ const cache = new ModelCache(options?.cacheName);
739
+ const buffer = await cache.getModel(modelUrl, options?.forceDownload, options);
740
+ return _EdgeInfer.fromBuffer(buffer, options);
741
+ }
742
+ /**
743
+ * Create an EdgeInfer instance from a pre-loaded model buffer.
744
+ */
745
+ static async fromBuffer(buffer, options) {
746
+ checkDownloadAbort(options?.signal);
747
+ const { session, provider } = await RuntimeManager.createSession(
748
+ buffer,
749
+ options?.executionProviders
750
+ );
751
+ try {
752
+ checkDownloadAbort(options?.signal);
753
+ let tokenizer;
754
+ if (options?.tokenizerUrl) {
755
+ tokenizer = await Tokenizer.fromUrl(options.tokenizerUrl, options);
756
+ } else if (options?.tokenizerConfig) {
757
+ tokenizer = new Tokenizer(options.tokenizerConfig);
758
+ }
759
+ checkDownloadAbort(options?.signal);
760
+ return new _EdgeInfer(session, buffer.byteLength, provider, options?.int64Inputs ?? false, tokenizer);
761
+ } catch (error) {
762
+ try {
763
+ await session.release();
764
+ } catch {
765
+ }
766
+ throw error;
767
+ }
768
+ }
769
+ /**
770
+ * Set or replace the tokenizer after model load.
771
+ */
772
+ setTokenizer(tokenizer) {
773
+ this.tokenizer = tokenizer;
774
+ }
775
+ /**
776
+ * Low-level tensor inference.
777
+ * Pass raw tensor inputs and get raw tensor outputs.
778
+ *
779
+ * `shapes` optionally gives the ONNX tensor shape per input name. When an
780
+ * input has no explicit shape it defaults to `[1, data.length]`, which is
781
+ * correct for 2-D sequence inputs (`input_ids`, `attention_mask`) but wrong
782
+ * for anything of higher rank — a vision model declaring
783
+ * `pixel_values: [1,3,224,224]` is rank 4 and rejects a rank-2 tensor with
784
+ * "Invalid rank for input". Pass `shapes` for such models.
785
+ */
786
+ async predict(inputs, shapes) {
787
+ const results = await this.runInference(inputs, shapes);
788
+ const outputMap = {};
789
+ for (const [key, tensor] of Object.entries(results)) {
790
+ outputMap[key] = tensor.data;
791
+ }
792
+ return outputMap;
793
+ }
794
+ /**
795
+ * Classify text using the loaded model.
796
+ * Requires a tokenizer to be configured.
797
+ */
798
+ async classify(text) {
799
+ this.requireTokenizer();
800
+ const encoding = this.tokenizer.encode(text, { addSpecialTokens: true });
801
+ const inputs = this.textInputs(encoding);
802
+ const result = await this.predict(inputs);
803
+ const output = result[this._outputNames[0]];
804
+ if (!output) return [];
805
+ const probs = this.softmax(output);
806
+ return Array.from(probs).map((score, idx) => ({
807
+ label: `Class_${idx}`,
808
+ score
809
+ })).sort((a, b) => b.score - a.score);
810
+ }
811
+ /**
812
+ * Generate text embeddings for semantic search / RAG.
813
+ * Returns a Float32Array of embedding dimensions.
814
+ */
815
+ async embed(text) {
816
+ this.requireTokenizer();
817
+ const encoding = this.tokenizer.encode(text, { addSpecialTokens: true });
818
+ const inputs = this.textInputs(encoding);
819
+ const result = await this.runInference(inputs);
820
+ const output = result[this._outputNames[0]];
821
+ if (!output) return new Float32Array();
822
+ if (output.type !== "float32") {
823
+ throw new Error("EdgeInfer: Embedding output must be a float32 tensor.");
824
+ }
825
+ return this.poolEmbedding(output.data, output.dims, encoding.attentionMask);
826
+ }
827
+ /**
828
+ * Generate Matryoshka-truncated embeddings at specified dimension.
829
+ * Supports variable-dimension MRL embeddings (64d → 256d → 768d).
830
+ */
831
+ async embedMatryoshka(text, dimension) {
832
+ const fullEmbedding = await this.embed(text);
833
+ if (dimension >= fullEmbedding.length) return fullEmbedding;
834
+ const truncated = fullEmbedding.slice(0, dimension);
835
+ return this.l2Normalize(truncated);
836
+ }
837
+ /**
838
+ * Greedy or temperature-sampled autoregressive generation.
839
+ *
840
+ * Decoder models that expose `present.*` outputs and `past_key_values.*`
841
+ * inputs reuse those tensors across steps (KV-cache). Models without cache
842
+ * inputs are re-run over the full sequence each step.
843
+ */
844
+ async generate(prompt, options = {}) {
845
+ this.requireTokenizer();
846
+ const tokenIds = [];
847
+ for await (const event of this.generateStream(prompt, options)) {
848
+ tokenIds.push(event.tokenId);
849
+ }
850
+ return {
851
+ text: this.tokenizer.decode(tokenIds),
852
+ tokenIds
853
+ };
854
+ }
855
+ /**
856
+ * Stream generated tokens. The decoder loop is identical to generate(), but
857
+ * each token is yielded as soon as it is chosen so a UI can paint before the
858
+ * full completion finishes.
859
+ */
860
+ async *generateStream(prompt, options = {}) {
861
+ this.requireTokenizer();
862
+ const maxNewTokens = Math.max(0, options.maxNewTokens ?? 64);
863
+ const temperature = options.temperature ?? 0;
864
+ const topK = options.topK ?? 0;
865
+ const tokenizer = this.tokenizer;
866
+ const eosTokenId = options.eosTokenId ?? tokenizer.getSpecialTokenId("</s>") ?? tokenizer.getSpecialTokenId("<|endoftext|>");
867
+ const encoding = tokenizer.encode(prompt, { addSpecialTokens: true });
868
+ const promptIds = Array.from(encoding.inputIds);
869
+ const int64 = this.useInt64Inputs();
870
+ const makeIds = (values) => int64 ? BigInt64Array.from(values, (v) => BigInt(v)) : Int32Array.from(values);
871
+ const dtype = int64 ? "int64" : "int32";
872
+ let current = promptIds;
873
+ let past = {};
874
+ const generated = [];
875
+ for (let step = 0; step < maxNewTokens; step++) {
876
+ throwIfAborted(options.signal);
877
+ const totalLength = promptIds.length + generated.length;
878
+ const inputs = {
879
+ input_ids: new ort2.Tensor(dtype, makeIds(current), [1, current.length]),
880
+ attention_mask: new ort2.Tensor(dtype, makeIds(new Array(totalLength).fill(1)), [1, totalLength])
881
+ };
882
+ if (this._inputNames.includes("position_ids")) {
883
+ const offset = totalLength - current.length;
884
+ inputs.position_ids = new ort2.Tensor(
885
+ dtype,
886
+ makeIds(current.map((_, i) => offset + i)),
887
+ [1, current.length]
888
+ );
889
+ }
890
+ for (const [name, tensor] of Object.entries(past)) {
891
+ if (this._inputNames.includes(name)) inputs[name] = tensor;
892
+ }
893
+ const results = await this.session.run(inputs);
894
+ const logits = this.findLogits(results);
895
+ const next = this.pickNextToken(logits, temperature, topK);
896
+ generated.push(next);
897
+ yield {
898
+ tokenId: next,
899
+ text: tokenizer.decode([next]),
900
+ index: generated.length - 1
901
+ };
902
+ if (eosTokenId !== void 0 && next === eosTokenId) break;
903
+ const nextPast = {};
904
+ for (const outputName of this._outputNames) {
905
+ const inputName = _EdgeInfer.pastInputName(outputName);
906
+ if (inputName && this._inputNames.includes(inputName) && results[outputName]) {
907
+ nextPast[inputName] = results[outputName];
908
+ }
909
+ }
910
+ if (Object.keys(nextPast).length > 0) {
911
+ past = nextPast;
912
+ current = [next];
913
+ } else {
914
+ past = {};
915
+ current = [...promptIds, ...generated];
916
+ }
917
+ }
918
+ }
919
+ /**
920
+ * Stream a chat completion using a role-marker template. The default template
921
+ * is model-agnostic; pass `chatTemplate` to match the target tokenizer.
922
+ */
923
+ async *chatStream(messages, options = {}) {
924
+ this.requireTokenizer();
925
+ const normalized = normalizeChatMessages(messages);
926
+ const prompt = (options.chatTemplate ?? defaultChatTemplate)(normalized);
927
+ yield* this.generateStream(prompt, options);
928
+ }
929
+ /** Collect a chat completion. */
930
+ async chat(messages, options = {}) {
931
+ this.requireTokenizer();
932
+ const normalized = normalizeChatMessages(messages);
933
+ const prompt = (options.chatTemplate ?? defaultChatTemplate)(normalized);
934
+ const tokenIds = [];
935
+ for await (const event of this.generateStream(prompt, options)) {
936
+ tokenIds.push(event.tokenId);
937
+ }
938
+ return {
939
+ text: this.tokenizer.decode(tokenIds),
940
+ tokenIds,
941
+ prompt
942
+ };
943
+ }
944
+ /** Pick the last-token logits from the first output that looks like logits. */
945
+ findLogits(results) {
946
+ const exact = results["logits"];
947
+ if (exact) return exact;
948
+ for (const name of Object.keys(results)) {
949
+ if (name.toLowerCase().includes("logits")) return results[name];
950
+ }
951
+ return results[this._outputNames[0]];
952
+ }
953
+ /** Argmax (temperature 0) or temperature/top-k sampling over last-token logits. */
954
+ pickNextToken(logits, temperature, topK) {
955
+ const dims = logits.dims;
956
+ const vocabSize = dims[dims.length - 1];
957
+ const seqLen = dims.length >= 3 ? dims[dims.length - 2] : 1;
958
+ const data = logits.data;
959
+ const offset = Math.max(0, (seqLen - 1) * vocabSize);
960
+ if (temperature <= 0) {
961
+ let best = 0;
962
+ let bestValue = -Infinity;
963
+ for (let i = 0; i < vocabSize; i++) {
964
+ const value = data[offset + i];
965
+ if (value > bestValue) {
966
+ bestValue = value;
967
+ best = i;
968
+ }
969
+ }
970
+ return best;
971
+ }
972
+ const scaled = new Float32Array(vocabSize);
973
+ let maxLogit = -Infinity;
974
+ for (let i = 0; i < vocabSize; i++) {
975
+ scaled[i] = data[offset + i] / temperature;
976
+ if (scaled[i] > maxLogit) maxLogit = scaled[i];
977
+ }
978
+ const candidates = [];
979
+ let sum = 0;
980
+ for (let i = 0; i < vocabSize; i++) {
981
+ const prob = Math.exp(scaled[i] - maxLogit);
982
+ candidates.push({ id: i, prob });
983
+ sum += prob;
984
+ }
985
+ if (topK > 0 && topK < candidates.length) {
986
+ candidates.sort((a, b) => b.prob - a.prob);
987
+ candidates.length = topK;
988
+ sum = candidates.reduce((acc, c) => acc + c.prob, 0);
989
+ }
990
+ let sample2 = Math.random() * sum;
991
+ for (const candidate of candidates) {
992
+ sample2 -= candidate.prob;
993
+ if (sample2 <= 0) return candidate.id;
994
+ }
995
+ return candidates[candidates.length - 1].id;
996
+ }
997
+ static pastInputName(outputName) {
998
+ if (!outputName.startsWith("present")) return null;
999
+ return outputName.replace(/^present/, "past_key_values");
1000
+ }
1001
+ useInt64Inputs() {
1002
+ if (this._int64Inputs) return true;
1003
+ const metadata = this.session.inputMetadata;
1004
+ return metadata?.["input_ids"]?.type === "int64";
1005
+ }
1006
+ /** Build input_ids/attention_mask tensors in the dtype the model declares. */
1007
+ textInputs(encoding) {
1008
+ const int64 = this.useInt64Inputs();
1009
+ const inputs = {
1010
+ input_ids: int64 ? encoding.inputIdsInt64 : encoding.inputIds,
1011
+ attention_mask: int64 ? encoding.attentionMaskInt64 : encoding.attentionMask
1012
+ };
1013
+ if (this._inputNames.includes("token_type_ids")) {
1014
+ inputs["token_type_ids"] = int64 ? new BigInt64Array(encoding.inputIds.length) : new Int32Array(encoding.inputIds.length);
1015
+ }
1016
+ return inputs;
1017
+ }
1018
+ /**
1019
+ * Sentiment analysis.
1020
+ */
1021
+ async sentiment(text) {
1022
+ const results = await this.classify(text);
1023
+ if (results.length === 0) return { label: "unknown", score: 0 };
1024
+ const best = results[0];
1025
+ const labelMap = {
1026
+ "Class_0": "negative",
1027
+ "Class_1": "positive",
1028
+ "Class_2": "neutral"
1029
+ };
1030
+ return {
1031
+ label: labelMap[best.label] || best.label,
1032
+ score: best.score
1033
+ };
1034
+ }
1035
+ /**
1036
+ * Classify an image using a vision model.
1037
+ */
1038
+ async classifyImage(imageData, layout) {
1039
+ const effectiveLayout = layout ?? "NCHW";
1040
+ const tensorData = ImageProcessor.imageDataToFloat32Array(imageData, { layout: effectiveLayout });
1041
+ const inputName = this._inputNames.find(
1042
+ (n) => n === "input" || n === "pixel_values" || n === "image"
1043
+ ) || this._inputNames[0];
1044
+ const result = await this.predict(
1045
+ { [inputName]: tensorData },
1046
+ { [inputName]: _EdgeInfer.imageShape(imageData, effectiveLayout) }
1047
+ );
1048
+ const output = result[this._outputNames[0]];
1049
+ if (!output) return [];
1050
+ const probs = this.softmax(output);
1051
+ return Array.from(probs).map((score, idx) => ({
1052
+ label: `Class_${idx}`,
1053
+ score
1054
+ })).sort((a, b) => b.score - a.score).slice(0, 10);
1055
+ }
1056
+ /**
1057
+ * Detect objects in an image.
1058
+ */
1059
+ async detectObjects(imageData) {
1060
+ const tensorData = ImageProcessor.imageDataToFloat32Array(imageData);
1061
+ const inputName = this._inputNames.find(
1062
+ (n) => n === "input" || n === "pixel_values" || n === "image"
1063
+ ) || this._inputNames[0];
1064
+ const result = await this.predict(
1065
+ { [inputName]: tensorData },
1066
+ { [inputName]: _EdgeInfer.imageShape(imageData, "NCHW") }
1067
+ );
1068
+ const output = result[this._outputNames[0]];
1069
+ const detections = [];
1070
+ if (output) {
1071
+ for (let i = 0; i < output.length; i += 6) {
1072
+ const confidence = output[i + 4];
1073
+ if (confidence > 0.5) {
1074
+ detections.push({
1075
+ bbox: [output[i], output[i + 1], output[i + 2], output[i + 3]],
1076
+ score: confidence,
1077
+ label: `Object_${Math.round(output[i + 5])}`
1078
+ });
1079
+ }
1080
+ }
1081
+ }
1082
+ return detections.sort((a, b) => b.score - a.score);
1083
+ }
1084
+ /**
1085
+ * Check GPU capabilities and recommended configuration.
1086
+ */
1087
+ static async getCapabilities() {
1088
+ return RuntimeManager.detectCapabilities();
1089
+ }
1090
+ /**
1091
+ * Check if WebGPU is available for accelerated inference.
1092
+ */
1093
+ static async isWebGPUAvailable() {
1094
+ const caps = await RuntimeManager.detectCapabilities();
1095
+ return caps.hasWebGPU;
1096
+ }
1097
+ get inputNames() {
1098
+ return this._inputNames;
1099
+ }
1100
+ get outputNames() {
1101
+ return this._outputNames;
1102
+ }
1103
+ get modelSize() {
1104
+ return this._modelSize;
1105
+ }
1106
+ get executionProvider() {
1107
+ return this._provider;
1108
+ }
1109
+ /**
1110
+ * Release all resources (ONNX session, GPU memory).
1111
+ */
1112
+ dispose() {
1113
+ if (this.session) {
1114
+ this.session.release();
1115
+ }
1116
+ this.tokenizer = null;
1117
+ RuntimeManager.resetCache();
1118
+ }
1119
+ // --- Private Helpers ---
1120
+ /** Run without discarding the output tensor dimensions needed by embed(). */
1121
+ async runInference(inputs, shapes) {
1122
+ const tensorInputs = {};
1123
+ for (const [key, data] of Object.entries(inputs)) {
1124
+ let type;
1125
+ if (data instanceof Float32Array) {
1126
+ type = "float32";
1127
+ } else if (data instanceof BigInt64Array) {
1128
+ type = "int64";
1129
+ } else {
1130
+ type = "int32";
1131
+ }
1132
+ const shape = shapes?.[key] ?? [1, data.length];
1133
+ tensorInputs[key] = new ort2.Tensor(type, data, shape);
1134
+ }
1135
+ return this.session.run(tensorInputs);
1136
+ }
1137
+ /** ONNX tensor shape for a 3-channel image tensor in the given layout. */
1138
+ static imageShape(imageData, layout) {
1139
+ const { width, height } = imageData;
1140
+ return layout === "NCHW" ? [1, 3, height, width] : [1, height, width, 3];
1141
+ }
1142
+ requireTokenizer() {
1143
+ if (!this.tokenizer) {
1144
+ throw new Error(
1145
+ "EdgeInfer: No tokenizer configured. Either pass tokenizerUrl/tokenizerConfig in options, or call setTokenizer() before using text-based inference methods."
1146
+ );
1147
+ }
1148
+ }
1149
+ softmax(logits) {
1150
+ let maxLogit = -Infinity;
1151
+ for (let i = 0; i < logits.length; i++) {
1152
+ if (logits[i] > maxLogit) maxLogit = logits[i];
1153
+ }
1154
+ const exps = new Float32Array(logits.length);
1155
+ let sumExp = 0;
1156
+ for (let i = 0; i < logits.length; i++) {
1157
+ exps[i] = Math.exp(logits[i] - maxLogit);
1158
+ sumExp += exps[i];
1159
+ }
1160
+ for (let i = 0; i < exps.length; i++) {
1161
+ exps[i] /= sumExp;
1162
+ }
1163
+ return exps;
1164
+ }
1165
+ poolEmbedding(embeddings, dims, attentionMask) {
1166
+ if ((dims.length === 1 || dims.length === 2 && dims[0] === 1) && dims[dims.length - 1] > 0) {
1167
+ return this.l2Normalize(embeddings);
1168
+ }
1169
+ if (dims.length !== 3 || dims[0] !== 1 || dims[1] < 1 || dims[2] < 1) {
1170
+ throw new Error(
1171
+ `EdgeInfer: Unsupported embedding output shape [${dims.join(", ")}]. Expected [hidden], [1, hidden], or [1, sequence, hidden].`
1172
+ );
1173
+ }
1174
+ const [, seqLen, hiddenDim] = dims;
1175
+ if (seqLen !== attentionMask.length) {
1176
+ throw new Error(
1177
+ `EdgeInfer: Embedding sequence length ${seqLen} does not match attention mask length ${attentionMask.length}.`
1178
+ );
1179
+ }
1180
+ const pooled = new Float32Array(hiddenDim);
1181
+ let maskSum = 0;
1182
+ for (let t = 0; t < seqLen; t++) {
1183
+ if (attentionMask[t] === 0) continue;
1184
+ maskSum++;
1185
+ for (let d = 0; d < hiddenDim; d++) {
1186
+ pooled[d] += embeddings[t * hiddenDim + d];
1187
+ }
1188
+ }
1189
+ if (maskSum > 0) {
1190
+ for (let d = 0; d < hiddenDim; d++) {
1191
+ pooled[d] /= maskSum;
1192
+ }
1193
+ }
1194
+ return this.l2Normalize(pooled);
1195
+ }
1196
+ l2Normalize(vec) {
1197
+ let norm = 0;
1198
+ for (let i = 0; i < vec.length; i++) {
1199
+ norm += vec[i] * vec[i];
1200
+ }
1201
+ norm = Math.sqrt(norm);
1202
+ if (norm === 0) return vec;
1203
+ const result = new Float32Array(vec.length);
1204
+ for (let i = 0; i < vec.length; i++) {
1205
+ result[i] = vec[i] / norm;
1206
+ }
1207
+ return result;
1208
+ }
1209
+ };
1210
+
1211
+ // src/matryoshka.ts
1212
+ function l2Normalize(vector, dimension = vector.length) {
1213
+ const output = new Float32Array(dimension);
1214
+ if (dimension > vector.length) {
1215
+ throw new Error(`Cannot normalize ${dimension} dims from a ${vector.length}-dim vector`);
1216
+ }
1217
+ let norm = 0;
1218
+ for (let i = 0; i < dimension; i++) norm += vector[i] * vector[i];
1219
+ norm = Math.sqrt(norm);
1220
+ if (norm === 0) return output;
1221
+ for (let i = 0; i < dimension; i++) output[i] = vector[i] / norm;
1222
+ return output;
1223
+ }
1224
+ function truncateAndNormalize(vector, dimension) {
1225
+ if (!Number.isInteger(dimension) || dimension < 1) {
1226
+ throw new Error("dimension must be a positive integer");
1227
+ }
1228
+ return l2Normalize(vector, dimension);
1229
+ }
1230
+ function quantizeInt8(vector, dimension = vector.length) {
1231
+ let maxAbs = 0;
1232
+ for (let i = 0; i < dimension; i++) {
1233
+ const value = Math.abs(vector[i]);
1234
+ if (value > maxAbs) maxAbs = value;
1235
+ }
1236
+ const scale = maxAbs > 0 ? maxAbs / 127 : 1;
1237
+ const data = new Int8Array(dimension);
1238
+ for (let i = 0; i < dimension; i++) {
1239
+ const value = Math.round(vector[i] / scale);
1240
+ data[i] = value > 127 ? 127 : value < -127 ? -127 : value;
1241
+ }
1242
+ return { data, scale };
1243
+ }
1244
+ function dequantizeInt8(data, scale) {
1245
+ const output = new Float32Array(data.length);
1246
+ for (let i = 0; i < data.length; i++) output[i] = data[i] * scale;
1247
+ return output;
1248
+ }
1249
+ function dot(a, b, dimension) {
1250
+ let sum = 0;
1251
+ for (let i = 0; i < dimension; i++) sum += a[i] * b[i];
1252
+ return sum;
1253
+ }
1254
+ var TopKHeap = class {
1255
+ constructor(k) {
1256
+ this.k = k;
1257
+ this.items = [];
1258
+ if (k < 1) throw new Error("top-k must be positive");
1259
+ }
1260
+ push(index, score) {
1261
+ if (this.items.length < this.k) {
1262
+ this.items.push({ index, score });
1263
+ if (this.items.length === this.k) this.heapify();
1264
+ return;
1265
+ }
1266
+ if (score > this.items[0].score) {
1267
+ this.items[0] = { index, score };
1268
+ this.siftDown(0);
1269
+ }
1270
+ }
1271
+ heapify() {
1272
+ for (let i = Math.floor(this.items.length / 2) - 1; i >= 0; i--) this.siftDown(i);
1273
+ }
1274
+ siftDown(start) {
1275
+ let index = start;
1276
+ const length = this.items.length;
1277
+ for (; ; ) {
1278
+ const left = index * 2 + 1;
1279
+ const right = left + 1;
1280
+ let smallest = index;
1281
+ if (left < length && this.items[left].score < this.items[smallest].score) smallest = left;
1282
+ if (right < length && this.items[right].score < this.items[smallest].score) smallest = right;
1283
+ if (smallest === index) break;
1284
+ [this.items[index], this.items[smallest]] = [this.items[smallest], this.items[index]];
1285
+ index = smallest;
1286
+ }
1287
+ }
1288
+ values() {
1289
+ return [...this.items].sort((a, b) => b.score - a.score);
1290
+ }
1291
+ };
1292
+ var MemoryVectorStore = class {
1293
+ constructor() {
1294
+ this.records = [];
1295
+ }
1296
+ async add(records) {
1297
+ this.records.push(...records);
1298
+ }
1299
+ async getAll() {
1300
+ return this.records.slice();
1301
+ }
1302
+ async clear() {
1303
+ this.records = [];
1304
+ }
1305
+ async count() {
1306
+ return this.records.length;
1307
+ }
1308
+ };
1309
+ var IndexedDbVectorStore = class {
1310
+ constructor(dbName = "edgeinfer-matryoshka", storeName = "vectors") {
1311
+ this.dbName = dbName;
1312
+ this.storeName = storeName;
1313
+ this.dbPromise = null;
1314
+ }
1315
+ open() {
1316
+ if (this.dbPromise) return this.dbPromise;
1317
+ if (typeof indexedDB === "undefined") {
1318
+ return Promise.reject(new Error("IndexedDB is unavailable in this environment"));
1319
+ }
1320
+ this.dbPromise = new Promise((resolve, reject) => {
1321
+ const request = indexedDB.open(this.dbName, 1);
1322
+ request.onupgradeneeded = () => {
1323
+ const db = request.result;
1324
+ if (!db.objectStoreNames.contains(this.storeName)) {
1325
+ db.createObjectStore(this.storeName, { keyPath: "id" });
1326
+ }
1327
+ };
1328
+ request.onsuccess = () => resolve(request.result);
1329
+ request.onerror = () => reject(request.error ?? new Error("IndexedDB open failed"));
1330
+ });
1331
+ return this.dbPromise;
1332
+ }
1333
+ async add(records) {
1334
+ const db = await this.open();
1335
+ await new Promise((resolve, reject) => {
1336
+ const transaction = db.transaction(this.storeName, "readwrite");
1337
+ const store = transaction.objectStore(this.storeName);
1338
+ for (const record of records) {
1339
+ const stored = {
1340
+ id: record.id,
1341
+ text: record.text,
1342
+ metadata: record.metadata,
1343
+ embedding: record.embedding.slice().buffer
1344
+ };
1345
+ store.put(stored);
1346
+ }
1347
+ transaction.oncomplete = () => resolve();
1348
+ transaction.onerror = () => reject(transaction.error ?? new Error("IndexedDB write failed"));
1349
+ transaction.onabort = () => reject(transaction.error ?? new Error("IndexedDB write aborted"));
1350
+ });
1351
+ }
1352
+ async getAll() {
1353
+ const db = await this.open();
1354
+ return new Promise((resolve, reject) => {
1355
+ const transaction = db.transaction(this.storeName, "readonly");
1356
+ const store = transaction.objectStore(this.storeName);
1357
+ const request = store.getAll();
1358
+ request.onsuccess = () => {
1359
+ const records = request.result.map((stored) => ({
1360
+ id: stored.id,
1361
+ text: stored.text,
1362
+ metadata: stored.metadata,
1363
+ embedding: new Float32Array(stored.embedding)
1364
+ }));
1365
+ resolve(records);
1366
+ };
1367
+ request.onerror = () => reject(request.error ?? new Error("IndexedDB read failed"));
1368
+ });
1369
+ }
1370
+ async clear() {
1371
+ const db = await this.open();
1372
+ await new Promise((resolve, reject) => {
1373
+ const transaction = db.transaction(this.storeName, "readwrite");
1374
+ transaction.objectStore(this.storeName).clear();
1375
+ transaction.oncomplete = () => resolve();
1376
+ transaction.onerror = () => reject(transaction.error ?? new Error("IndexedDB clear failed"));
1377
+ });
1378
+ }
1379
+ async count() {
1380
+ const db = await this.open();
1381
+ return new Promise((resolve, reject) => {
1382
+ const transaction = db.transaction(this.storeName, "readonly");
1383
+ const request = transaction.objectStore(this.storeName).count();
1384
+ request.onsuccess = () => resolve(request.result);
1385
+ request.onerror = () => reject(request.error ?? new Error("IndexedDB count failed"));
1386
+ });
1387
+ }
1388
+ close() {
1389
+ void this.dbPromise?.then((db) => db.close());
1390
+ this.dbPromise = null;
1391
+ }
1392
+ };
1393
+ var MatryoshkaIndex = class _MatryoshkaIndex {
1394
+ constructor(options = {}) {
1395
+ this.ids = [];
1396
+ this.texts = [];
1397
+ this.metadata = [];
1398
+ this.full = new Float32Array(0);
1399
+ this.coarse = new Int8Array(0);
1400
+ this.scales = new Float32Array(0);
1401
+ this.size = 0;
1402
+ this.dimensions = options.dimensions ?? 768;
1403
+ this.coarseDimension = options.coarseDimension ?? 64;
1404
+ this.refineTopK = options.refineTopK ?? 50;
1405
+ this.store = options.store;
1406
+ if (this.coarseDimension > this.dimensions) {
1407
+ throw new Error("coarseDimension cannot exceed dimensions");
1408
+ }
1409
+ }
1410
+ get count() {
1411
+ return this.size;
1412
+ }
1413
+ async build(records, { persist = true } = {}) {
1414
+ if (records.length === 0) throw new Error("At least one record is required");
1415
+ for (const record of records) {
1416
+ if (record.embedding.length !== this.dimensions) {
1417
+ throw new Error(
1418
+ `Record ${record.id} has ${record.embedding.length} dims; expected ${this.dimensions}`
1419
+ );
1420
+ }
1421
+ }
1422
+ this.ids = new Array(records.length);
1423
+ this.texts = new Array(records.length);
1424
+ this.metadata = new Array(records.length);
1425
+ this.full = new Float32Array(records.length * this.dimensions);
1426
+ this.coarse = new Int8Array(records.length * this.coarseDimension);
1427
+ this.scales = new Float32Array(records.length);
1428
+ this.size = records.length;
1429
+ for (let row = 0; row < records.length; row++) {
1430
+ const record = records[row];
1431
+ const normalized = l2Normalize(record.embedding, this.dimensions);
1432
+ this.ids[row] = record.id;
1433
+ this.texts[row] = record.text;
1434
+ this.metadata[row] = record.metadata;
1435
+ this.full.set(normalized, row * this.dimensions);
1436
+ const quantized = quantizeInt8(normalized, this.coarseDimension);
1437
+ this.coarse.set(quantized.data, row * this.coarseDimension);
1438
+ this.scales[row] = quantized.scale;
1439
+ }
1440
+ if (persist && this.store) {
1441
+ await this.store.clear();
1442
+ const normalizedRecords = records.map((record, row) => ({
1443
+ ...record,
1444
+ embedding: this.full.slice(row * this.dimensions, (row + 1) * this.dimensions)
1445
+ }));
1446
+ await this.store.add(normalizedRecords);
1447
+ }
1448
+ }
1449
+ static async load(store, options = {}) {
1450
+ const index = new _MatryoshkaIndex({ ...options, store });
1451
+ const records = await store.getAll();
1452
+ await index.build(records, { persist: false });
1453
+ return index;
1454
+ }
1455
+ search(query, options = {}) {
1456
+ if (this.size === 0) throw new Error("Index is empty; call build() first");
1457
+ const topK = options.topK ?? 10;
1458
+ const refine = options.refine ?? true;
1459
+ const coarseDimension = options.coarseDimension ?? this.coarseDimension;
1460
+ if (coarseDimension > this.coarseDimension) {
1461
+ throw new Error(
1462
+ `coarseDimension ${coarseDimension} exceeds the indexed stride ${this.coarseDimension}`
1463
+ );
1464
+ }
1465
+ const queryFull = l2Normalize(query, this.dimensions);
1466
+ const queryCoarse = quantizeInt8(queryFull, coarseDimension);
1467
+ const refineTopK = Math.max(
1468
+ topK,
1469
+ Math.min(this.refineTopK, options.refineTopK ?? this.refineTopK, this.size)
1470
+ );
1471
+ const heap = new TopKHeap(refineTopK);
1472
+ for (let row = 0; row < this.size; row++) {
1473
+ const offset = row * this.coarseDimension;
1474
+ let sum = 0;
1475
+ for (let i = 0; i < coarseDimension; i++) {
1476
+ sum += queryCoarse.data[i] * this.coarse[offset + i];
1477
+ }
1478
+ heap.push(row, sum * this.scales[row]);
1479
+ }
1480
+ let candidates = heap.values().map((candidate) => ({ ...candidate, coarseScore: candidate.score }));
1481
+ if (refine) {
1482
+ candidates = candidates.map((candidate) => ({
1483
+ ...candidate,
1484
+ score: dot(
1485
+ queryFull,
1486
+ this.full.subarray(
1487
+ candidate.index * this.dimensions,
1488
+ (candidate.index + 1) * this.dimensions
1489
+ ),
1490
+ this.dimensions
1491
+ )
1492
+ })).sort((a, b) => b.score - a.score);
1493
+ }
1494
+ return candidates.slice(0, topK).map((candidate) => ({
1495
+ id: this.ids[candidate.index],
1496
+ text: this.texts[candidate.index],
1497
+ metadata: this.metadata[candidate.index],
1498
+ score: candidate.score,
1499
+ coarseScore: candidate.coarseScore
1500
+ }));
1501
+ }
1502
+ exactSearchAtDimension(query, dimension, topK = 10) {
1503
+ if (dimension > this.dimensions) throw new Error("dimension exceeds index dimensions");
1504
+ const queryNormalized = truncateAndNormalize(query, dimension);
1505
+ const scored = [];
1506
+ for (let row = 0; row < this.size; row++) {
1507
+ const vector = truncateAndNormalize(
1508
+ this.full.subarray(row * this.dimensions, (row + 1) * this.dimensions),
1509
+ dimension
1510
+ );
1511
+ scored.push({ index: row, score: dot(queryNormalized, vector, dimension) });
1512
+ }
1513
+ return scored.sort((a, b) => b.score - a.score).slice(0, topK);
1514
+ }
1515
+ recallAtK(queries, groundTruth, dimension, topK = 10) {
1516
+ let hits = 0;
1517
+ let total = 0;
1518
+ for (let q = 0; q < queries.length; q++) {
1519
+ const expected = new Set(groundTruth[q].slice(0, topK));
1520
+ const candidates = this.exactSearchAtDimension(queries[q], dimension, topK);
1521
+ for (const candidate of candidates) {
1522
+ if (expected.has(this.ids[candidate.index])) hits++;
1523
+ }
1524
+ total += Math.min(topK, expected.size);
1525
+ }
1526
+ return total === 0 ? 0 : hits / total;
1527
+ }
1528
+ get idList() {
1529
+ return this.ids;
1530
+ }
1531
+ };
1532
+ async function embedMatryoshka(text, dimension, embed) {
1533
+ const full = await embed(text);
1534
+ return truncateAndNormalize(full, dimension);
1535
+ }
1536
+ function createMatryoshkaEmbedder(embed) {
1537
+ return (text, dimension = 768) => embedMatryoshka(text, dimension, embed);
1538
+ }
1539
+
1540
+ // src/speculative.ts
1541
+ var SpeculativeDecoder = class {
1542
+ constructor(options) {
1543
+ if (options.draft.vocabSize !== options.target.vocabSize) {
1544
+ throw new Error("Draft and target models must share a vocabulary size");
1545
+ }
1546
+ const gamma = options.gamma ?? 5;
1547
+ if (!Number.isInteger(gamma) || gamma < 1) {
1548
+ throw new Error("gamma must be a positive integer");
1549
+ }
1550
+ this.draft = options.draft;
1551
+ this.target = options.target;
1552
+ this.gamma = gamma;
1553
+ this.rng = options.rng ?? Math.random;
1554
+ this.eosTokenId = options.eosTokenId;
1555
+ }
1556
+ async generate(prompt, maxNewTokens = 64, { temperature = 1 } = {}) {
1557
+ if (prompt.length === 0) throw new Error("prompt must not be empty");
1558
+ if (maxNewTokens < 1) throw new Error("maxNewTokens must be positive");
1559
+ if (temperature <= 0) throw new Error("temperature must be positive");
1560
+ const started = performance.now();
1561
+ let draftState = await this.draft.prefill(prompt);
1562
+ let targetState = await this.target.prefill(prompt);
1563
+ const output = [];
1564
+ let draftTokens = 0;
1565
+ let acceptedTokens = 0;
1566
+ let bonusTokens = 0;
1567
+ let rejectedTokens = 0;
1568
+ let targetCalls = 0;
1569
+ let draftCalls = 0;
1570
+ while (output.length < maxNewTokens) {
1571
+ const gamma = Math.min(this.gamma, maxNewTokens - output.length);
1572
+ const draft = await this.draftBatch(draftState, gamma, temperature, draftCalls);
1573
+ draftState = draft.state;
1574
+ draftCalls = draft.calls;
1575
+ draftTokens += draft.tokens.length;
1576
+ const prefixLength = targetState.tokens.length;
1577
+ const verification = await this.target.decode(targetState, draft.tokens);
1578
+ targetCalls++;
1579
+ targetState = verification.state;
1580
+ let rejected = false;
1581
+ for (let i = 0; i < draft.tokens.length; i++) {
1582
+ const token = draft.tokens[i];
1583
+ const q = draft.probabilities[i][token];
1584
+ const p = probability(verification.logits[i], token, temperature);
1585
+ if (this.rng() < Math.min(1, p / q)) {
1586
+ output.push(token);
1587
+ acceptedTokens++;
1588
+ if (token === this.eosTokenId || output.length >= maxNewTokens) break;
1589
+ continue;
1590
+ }
1591
+ const residual = residualDistribution(
1592
+ verification.logits[i],
1593
+ draft.probabilities[i],
1594
+ temperature
1595
+ );
1596
+ const replacement = sample(residual, this.rng());
1597
+ output.push(replacement);
1598
+ rejectedTokens++;
1599
+ rejected = true;
1600
+ const keep = prefixLength + i;
1601
+ targetState = this.target.truncate(targetState, keep);
1602
+ draftState = this.draft.truncate(draftState, keep);
1603
+ const targetAdvance2 = await this.target.decode(targetState, [replacement]);
1604
+ targetCalls++;
1605
+ targetState = targetAdvance2.state;
1606
+ const draftAdvance2 = await this.draft.decode(draftState, [replacement]);
1607
+ draftCalls++;
1608
+ draftState = draftAdvance2.state;
1609
+ break;
1610
+ }
1611
+ if (rejected) continue;
1612
+ if (output.length >= maxNewTokens || output[output.length - 1] === this.eosTokenId) break;
1613
+ const bonusDistribution = softmax(
1614
+ verification.logits[draft.tokens.length],
1615
+ temperature
1616
+ );
1617
+ const bonus = sample(bonusDistribution, this.rng());
1618
+ output.push(bonus);
1619
+ bonusTokens++;
1620
+ const targetAdvance = await this.target.decode(targetState, [bonus]);
1621
+ targetCalls++;
1622
+ targetState = targetAdvance.state;
1623
+ const draftAdvance = await this.draft.decode(draftState, [bonus]);
1624
+ draftCalls++;
1625
+ draftState = draftAdvance.state;
1626
+ if (bonus === this.eosTokenId) break;
1627
+ }
1628
+ const elapsedMs = performance.now() - started;
1629
+ const stats = {
1630
+ outputTokens: output.length,
1631
+ draftTokens,
1632
+ acceptedTokens,
1633
+ bonusTokens,
1634
+ rejectedTokens,
1635
+ targetCalls,
1636
+ draftCalls,
1637
+ acceptanceRate: draftTokens === 0 ? 0 : acceptedTokens / draftTokens,
1638
+ targetCallsPerToken: output.length === 0 ? 0 : targetCalls / output.length,
1639
+ elapsedMs,
1640
+ tokensPerSecond: elapsedMs === 0 ? 0 : output.length / elapsedMs * 1e3
1641
+ };
1642
+ return { tokens: output, stats };
1643
+ }
1644
+ async draftBatch(state, gamma, temperature, calls) {
1645
+ let current = state;
1646
+ const tokens = [];
1647
+ const probabilities = [];
1648
+ const first = await this.draft.decode(current, []);
1649
+ calls++;
1650
+ let nextDistribution = first.logits[0];
1651
+ for (let i = 0; i < gamma; i++) {
1652
+ const distribution = softmax(nextDistribution, temperature);
1653
+ const token = sample(distribution, this.rng());
1654
+ tokens.push(token);
1655
+ probabilities.push(distribution);
1656
+ const advanced = await this.draft.decode(current, [token]);
1657
+ calls++;
1658
+ current = advanced.state;
1659
+ nextDistribution = advanced.logits[advanced.logits.length - 1];
1660
+ }
1661
+ return { state: current, tokens, probabilities, calls };
1662
+ }
1663
+ };
1664
+ function softmax(logits, temperature = 1) {
1665
+ let max = -Infinity;
1666
+ for (let i = 0; i < logits.length; i++) max = Math.max(max, logits[i]);
1667
+ const output = new Float32Array(logits.length);
1668
+ let sum = 0;
1669
+ for (let i = 0; i < logits.length; i++) {
1670
+ const value = Math.exp((logits[i] - max) / temperature);
1671
+ output[i] = value;
1672
+ sum += value;
1673
+ }
1674
+ if (sum === 0) return output;
1675
+ for (let i = 0; i < output.length; i++) output[i] /= sum;
1676
+ return output;
1677
+ }
1678
+ function probability(logits, token, temperature = 1) {
1679
+ return softmax(logits, temperature)[token] ?? 0;
1680
+ }
1681
+ function sample(distribution, random) {
1682
+ let cumulative = 0;
1683
+ for (let i = 0; i < distribution.length; i++) {
1684
+ cumulative += distribution[i];
1685
+ if (random <= cumulative) return i;
1686
+ }
1687
+ return distribution.length - 1;
1688
+ }
1689
+ function residualDistribution(targetLogits, draftProbabilities, temperature = 1) {
1690
+ const p = softmax(targetLogits, temperature);
1691
+ const q = draftProbabilities;
1692
+ const residual = new Float32Array(p.length);
1693
+ let sum = 0;
1694
+ for (let i = 0; i < p.length; i++) {
1695
+ const value = Math.max(0, p[i] - q[i]);
1696
+ residual[i] = value;
1697
+ sum += value;
1698
+ }
1699
+ if (sum === 0) return p;
1700
+ for (let i = 0; i < residual.length; i++) residual[i] /= sum;
1701
+ return residual;
1702
+ }
1703
+ var MockLanguageModel = class {
1704
+ constructor(policy, vocabSize, latencyMs = 0) {
1705
+ this.policy = policy;
1706
+ this.vocabSize = vocabSize;
1707
+ this.latencyMs = latencyMs;
1708
+ }
1709
+ async prefill(tokens) {
1710
+ await delay(this.latencyMs);
1711
+ return { tokens: tokens.slice() };
1712
+ }
1713
+ async decode(state, tokens) {
1714
+ await delay(this.latencyMs);
1715
+ const logits = [];
1716
+ for (let i = 0; i <= tokens.length; i++) {
1717
+ logits.push(this.policy([...state.tokens, ...tokens.slice(0, i)]));
1718
+ }
1719
+ return { state: { tokens: [...state.tokens, ...tokens] }, logits };
1720
+ }
1721
+ truncate(state, length) {
1722
+ return { ...state, tokens: state.tokens.slice(0, length), cache: void 0 };
1723
+ }
1724
+ };
1725
+ function delay(ms) {
1726
+ if (ms <= 0) return Promise.resolve();
1727
+ return new Promise((resolve) => setTimeout(resolve, ms));
1728
+ }
1729
+
1730
+ // src/onnx-lm.ts
1731
+ var ort3 = __toESM(require("onnxruntime-web"));
1732
+ var OnnxCausalLM = class {
1733
+ constructor(session, vocabSize, options = {}) {
1734
+ this.session = session;
1735
+ if (!Number.isInteger(vocabSize) || vocabSize < 1) {
1736
+ throw new Error("vocabSize must be a positive integer");
1737
+ }
1738
+ this.vocabSize = vocabSize;
1739
+ this.inputIdsName = options.inputIdsName ?? "input_ids";
1740
+ this.attentionMaskName = options.attentionMaskName ?? "attention_mask";
1741
+ this.positionIdsName = options.positionIdsName;
1742
+ this.logitsName = options.logitsName ?? "logits";
1743
+ }
1744
+ async prefill(tokens) {
1745
+ if (tokens.length === 0) throw new Error("prefill requires a non-empty prompt");
1746
+ return { tokens: tokens.slice() };
1747
+ }
1748
+ async decode(state, tokens) {
1749
+ const all = [...state.tokens, ...tokens];
1750
+ if (all.length === 0) throw new Error("decode requires at least one token");
1751
+ if (state.tokens.length === 0) throw new Error("decode requires a prefilled state");
1752
+ const feeds = {
1753
+ [this.inputIdsName]: new ort3.Tensor(
1754
+ "int64",
1755
+ BigInt64Array.from(all, (token) => BigInt(token)),
1756
+ [1, all.length]
1757
+ ),
1758
+ [this.attentionMaskName]: new ort3.Tensor(
1759
+ "int64",
1760
+ BigInt64Array.from({ length: all.length }, () => 1n),
1761
+ [1, all.length]
1762
+ )
1763
+ };
1764
+ if (this.positionIdsName) {
1765
+ feeds[this.positionIdsName] = new ort3.Tensor(
1766
+ "int64",
1767
+ BigInt64Array.from(all, (_, index) => BigInt(index)),
1768
+ [1, all.length]
1769
+ );
1770
+ }
1771
+ const outputs = await this.session.run(feeds);
1772
+ const tensor = outputs[this.logitsName];
1773
+ if (!tensor) {
1774
+ throw new Error(`ONNX model did not return "${this.logitsName}"`);
1775
+ }
1776
+ const dims = tensor.dims;
1777
+ const vocab = dims[dims.length - 1];
1778
+ if (vocab !== this.vocabSize) {
1779
+ throw new Error(
1780
+ `ONNX logits vocabulary ${vocab} does not match configured ${this.vocabSize}`
1781
+ );
1782
+ }
1783
+ const data = tensor.data;
1784
+ const seq = all.length;
1785
+ const expected = seq * vocab;
1786
+ if (data.length < expected) {
1787
+ throw new Error(`ONNX logits tensor is too small (${data.length} < ${expected})`);
1788
+ }
1789
+ const start = state.tokens.length - 1;
1790
+ const logits = [];
1791
+ for (let i = 0; i <= tokens.length; i++) {
1792
+ const position = start + i;
1793
+ const offset = dims.length === 3 ? position * vocab : position * vocab;
1794
+ logits.push(new Float32Array(data.subarray(offset, offset + vocab)));
1795
+ }
1796
+ return { state: { tokens: all }, logits };
1797
+ }
1798
+ truncate(state, length) {
1799
+ return { tokens: state.tokens.slice(0, length) };
1800
+ }
1801
+ };
1802
+ // Annotate the CommonJS export names for ESM import in node:
1803
+ 0 && (module.exports = {
1804
+ EdgeInfer,
1805
+ EventEmitter,
1806
+ ImageProcessor,
1807
+ IndexedDbVectorStore,
1808
+ MatryoshkaIndex,
1809
+ MemoryVectorStore,
1810
+ MockLanguageModel,
1811
+ ModelCache,
1812
+ OnnxCausalLM,
1813
+ RuntimeManager,
1814
+ SpeculativeDecoder,
1815
+ Tokenizer,
1816
+ createMatryoshkaEmbedder,
1817
+ dequantizeInt8,
1818
+ embedMatryoshka,
1819
+ l2Normalize,
1820
+ probability,
1821
+ quantizeInt8,
1822
+ residualDistribution,
1823
+ sample,
1824
+ softmax,
1825
+ truncateAndNormalize
1826
+ });
1827
+ //# sourceMappingURL=index.js.map