Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 8x 8x 8x 8x 8x 8x 8x 9x 9x 1x 8x 8x 8x 8x 8x 8x 8x 8x 8x 8x 8x 8x 7x 8x 7x 7x 7x 7x 7x 7x 7x 2x 2x 5x 2x 8x 8x 7x | /**
* NodeEmbedder — ONNX Runtime Node.js code embedding.
*
* Implements IEmbedder using onnxruntime-node (AVX2/AVX-512 native).
*/
import type {
IEmbedder,
BatchOptions,
ModelInfo,
TokenizedInput,
} from '@agentix-e/embed-code-core';
import {
WordPieceTokenizer,
meanPool,
l2Normalize,
processBatch,
} from '@agentix-e/embed-code-core';
import { NodeOrtBackend } from './ort-backend.js';
import * as path from 'node:path';
import * as fs from 'node:fs';
export interface NodeEmbedderOptions {
modelPath: string;
tokenizerPath?: string;
}
export class NodeEmbedder implements IEmbedder {
private readonly backend = new NodeOrtBackend();
private session: Awaited<ReturnType<NodeOrtBackend['createSession']>> | null = null;
readonly dimensions: number;
readonly maxSequenceLength: number;
readonly modelInfo: ModelInfo;
private tokenizer: WordPieceTokenizer;
private _modelPath: string;
private constructor(
options: NodeEmbedderOptions,
tokenizer: WordPieceTokenizer,
dimensions: number,
maxSequenceLength: number,
) {
this._modelPath = options.modelPath;
this.tokenizer = tokenizer;
this.dimensions = dimensions;
this.maxSequenceLength = maxSequenceLength;
this.modelInfo = {
name: 'nomic-embed-code',
version: 'v1.5',
dimensions,
maxSequenceLength,
vocabSize: tokenizer.vocabSize,
quantization: 'int8',
};
}
/** Create from a model file. The tokenizer.json is expected alongside the model. */
static async create(options: NodeEmbedderOptions): Promise<NodeEmbedder> {
const tokenizerPath =
options.tokenizerPath ?? path.join(path.dirname(options.modelPath), 'tokenizer.json');
if (!fs.existsSync(tokenizerPath)) {
throw new Error(`Tokenizer not found: ${tokenizerPath}`);
}
const descriptorPath = path.join(path.dirname(options.modelPath), 'model-descriptor.json');
let dimensions = 768;
let maxSequenceLength = 512;
Eif (fs.existsSync(descriptorPath)) {
try {
const raw = fs.readFileSync(descriptorPath, 'utf-8');
const descriptor = JSON.parse(raw);
dimensions = descriptor.architecture?.embedding_dim ?? dimensions;
maxSequenceLength = Math.min(
descriptor.tokenizer?.max_length ?? maxSequenceLength,
maxSequenceLength,
);
} catch {
// Use defaults
}
}
const tokenizer = WordPieceTokenizer.fromFile(tokenizerPath, maxSequenceLength);
const embedder = new NodeEmbedder(options, tokenizer, dimensions, maxSequenceLength);
embedder.session = await embedder.backend.createSession(options.modelPath);
return embedder;
}
/** Create from the embedded model in this package's models/ directory. */
static async createFromPackage(): Promise<NodeEmbedder> {
const __dirname = path.dirname(new URL(import.meta.url).pathname);
const modelPath = path.join(__dirname, '..', 'models', 'nomic-embed-code-v1.5.int8.onnx');
return NodeEmbedder.create({ modelPath });
}
async embed(text: string): Promise<Float32Array> {
if (!this.session) throw new Error('Session not initialized');
const tokens = this.tokenizer.tokenize(text);
const feeds = this.buildFeeds(tokens, 1);
const outputs = await this.session.run(feeds);
const hidden = outputs.last_hidden_state.data as Float32Array;
// Mean pool + L2 normalize
const pooled = meanPool(
hidden,
tokens.attentionMask,
1,
this.maxSequenceLength,
this.dimensions,
);
l2Normalize(pooled, 1, this.dimensions);
return pooled;
}
async embedBatch(texts: string[], options?: BatchOptions): Promise<Float32Array[]> {
const results: Float32Array[] = new Array(texts.length);
await processBatch(
texts,
async (text, index) => {
results[index] = await this.embed(text);
},
options,
);
return results;
}
async dispose(): Promise<void> {
this.session?.release();
this.session = null;
}
private buildFeeds(tokens: TokenizedInput, batch: number) {
return {
input_ids: this.backend.createTensor('int64', tokens.inputIds, [
batch,
this.maxSequenceLength,
]),
attention_mask: this.backend.createTensor('int64', tokens.attentionMask, [
batch,
this.maxSequenceLength,
]),
token_type_ids: this.backend.createTensor('int64', tokens.tokenTypeIds, [
batch,
this.maxSequenceLength,
]),
};
}
}
|