vectra 0.12.2 → 0.14.0

package/src/TransformersEmbeddings.spec.ts

@@ -0,0 +1,188 @@
+import { strict as assert } from 'node:assert';
+import { describe, it, beforeEach, afterEach } from 'mocha';
+import sinon from 'sinon';
+import { EmbeddingsModel } from './types';
+import * as transformersModule from '@huggingface/transformers';
+
+describe('TransformersEmbeddings', () => {
+    let TransformersEmbeddings: any;
+    let mockExtractor: sinon.SinonStub;
+    let mockTokenizer: any;
+    let sandbox: sinon.SinonSandbox;
+    let pipelineStub: sinon.SinonStub;
+
+    beforeEach(async () => {
+        sandbox = sinon.createSandbox();
+
+        // Create mock tokenizer
+        mockTokenizer = {
+            __call__: sandbox.stub().returns({
+                input_ids: { data: BigInt64Array.from([BigInt(1), BigInt(2), BigInt(3)]) }
+            }),
+            decode: sandbox.stub().returns('decoded text')
+        };
+        // Make it callable
+        const callableTokenizer = Object.assign(
+            (...args: any[]) => mockTokenizer.__call__(...args),
+            mockTokenizer
+        );
+
+        // Create mock extractor (feature extraction pipeline)
+        mockExtractor = sandbox.stub().callsFake(async (inputs: string | string[]) => {
+            const inputArray = Array.isArray(inputs) ? inputs : [inputs];
+            const batchSize = inputArray.length;
+            const embeddingDim = 4;
+
+            const data = new Float32Array(batchSize * embeddingDim);
+            for (let i = 0; i < batchSize; i++) {
+                data[i * embeddingDim] = 0.1;
+                data[i * embeddingDim + 1] = 0.2;
+                data[i * embeddingDim + 2] = 0.3;
+                data[i * embeddingDim + 3] = 0.4;
+            }
+
+            return {
+                data: data,
+                dims: [batchSize, embeddingDim]
+            };
+        });
+
+        // Attach tokenizer to the mock extractor so pipeline result has .tokenizer
+        (mockExtractor as any).tokenizer = callableTokenizer;
+
+        // Stub the pipeline function from @huggingface/transformers
+        pipelineStub = sandbox.stub(transformersModule, 'pipeline' as any).resolves(mockExtractor);
+
+        // Import TransformersEmbeddings fresh (uses the stubbed pipeline via dynamic import)
+        const mod = await import('./TransformersEmbeddings');
+        TransformersEmbeddings = mod.TransformersEmbeddings;
+    });
+
+    afterEach(() => {
+        sandbox.restore();
+    });
+
+    describe('create()', () => {
+        it('creates instance with default options', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+
+            assert.equal(embeddings.maxTokens, 512, 'default maxTokens should be 512');
+            assert.equal(embeddings.model, 'Xenova/all-MiniLM-L6-v2', 'default model should be all-MiniLM-L6-v2');
+
+            // Verify pipeline was called with correct arguments
+            assert.ok(pipelineStub.calledOnce, 'pipeline should be called once');
+            assert.equal(pipelineStub.firstCall.args[0], 'feature-extraction');
+            assert.equal(pipelineStub.firstCall.args[1], 'Xenova/all-MiniLM-L6-v2');
+        });
+
+        it('creates instance with custom options', async () => {
+            const embeddings = await TransformersEmbeddings.create({
+                model: 'Xenova/bge-small-en-v1.5',
+                maxTokens: 256,
+                device: 'cpu',
+                normalize: false,
+                pooling: 'cls'
+            });
+
+            assert.equal(embeddings.maxTokens, 256);
+            assert.equal(embeddings.model, 'Xenova/bge-small-en-v1.5');
+        });
+
+        it('implements EmbeddingsModel interface', async () => {
+            const embeddings: EmbeddingsModel = await TransformersEmbeddings.create();
+
+            assert.equal(typeof embeddings.maxTokens, 'number');
+            assert.equal(typeof embeddings.createEmbeddings, 'function');
+        });
+    });
+
+    describe('createEmbeddings()', () => {
+        it('generates embeddings for single string', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const result = await embeddings.createEmbeddings('hello world');
+
+            assert.equal(result.status, 'success');
+            assert.ok(result.output, 'output should be defined');
+            assert.equal(result.output!.length, 1, 'should have one embedding');
+            assert.equal(result.output![0].length, 4, 'embedding should have 4 dimensions');
+            const expected = [0.1, 0.2, 0.3, 0.4];
+            result.output![0].forEach((val: number, i: number) => {
+                assert.ok(Math.abs(val - expected[i]) < 0.001, `value ${val} should be close to ${expected[i]}`);
+            });
+            assert.equal(result.model, 'Xenova/all-MiniLM-L6-v2');
+        });
+
+        it('generates embeddings for string array', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const result = await embeddings.createEmbeddings(['hello', 'world']);
+
+            assert.equal(result.status, 'success');
+            assert.ok(result.output, 'output should be defined');
+            assert.equal(result.output!.length, 2, 'should have two embeddings');
+
+            assert.equal(mockExtractor.callCount, 1);
+            assert.deepEqual(mockExtractor.firstCall.args[0], ['hello', 'world']);
+        });
+
+        it('passes pooling and normalize options to extractor', async () => {
+            const embeddings = await TransformersEmbeddings.create({
+                pooling: 'cls',
+                normalize: false
+            });
+            await embeddings.createEmbeddings('test');
+
+            assert.ok(mockExtractor.calledOnce);
+            const options = mockExtractor.firstCall.args[1];
+            assert.equal(options.pooling, 'cls');
+            assert.equal(options.normalize, false);
+        });
+
+        it('returns error status on failure', async () => {
+            mockExtractor.rejects(new Error('Model inference failed'));
+
+            const embeddings = await TransformersEmbeddings.create();
+            const result = await embeddings.createEmbeddings('test');
+
+            assert.equal(result.status, 'error');
+            assert.ok(result.message?.includes('Model inference failed'));
+        });
+
+        it('handles empty string input', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const result = await embeddings.createEmbeddings('');
+
+            assert.equal(result.status, 'success');
+            assert.ok(result.output);
+            assert.equal(result.output!.length, 1);
+        });
+
+        it('handles empty array input', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const result = await embeddings.createEmbeddings([]);
+
+            assert.equal(result.status, 'success');
+            assert.ok(result.output);
+            assert.equal(result.output!.length, 0);
+        });
+    });
+
+    describe('getTokenizer()', () => {
+        it('returns a TransformersTokenizer instance', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const tokenizer = embeddings.getTokenizer();
+
+            assert.ok(tokenizer, 'tokenizer should be defined');
+            assert.equal(typeof tokenizer.encode, 'function');
+            assert.equal(typeof tokenizer.decode, 'function');
+        });
+
+        it('returns consistent tokenizer across calls', async () => {
+            const embeddings = await TransformersEmbeddings.create();
+            const tokenizer1 = embeddings.getTokenizer();
+            const tokenizer2 = embeddings.getTokenizer();
+
+            assert.ok(tokenizer1);
+            assert.ok(tokenizer2);
+        });
+    });
+});

package/src/TransformersEmbeddings.ts

@@ -0,0 +1,232 @@
+import { EmbeddingsModel, EmbeddingsResponse } from "./types";
+import { TransformersTokenizer } from "./TransformersTokenizer";
+import { FeatureExtractionPipeline, PreTrainedTokenizer } from "@huggingface/transformers";
+
+
+const DEFAULT_MODEL = 'Xenova/all-MiniLM-L6-v2';
+
+/**
+ * Type definition for the Transformers.js library.
+ * Used for dynamic import and type safety.
+ */
+type TransformersLibrary = typeof import('@huggingface/transformers');
+
+/**
+ * Configuration options for TransformersEmbeddings.
+ */
+export interface TransformersEmbeddingsOptions {
+    /**
+     * Optional. Model name/path to use for embeddings.
+     * @remarks
+     * Common models:
+     * - 'Xenova/all-MiniLM-L6-v2' (384 dimensions, fast, good quality)
+     * - 'Xenova/bge-small-en-v1.5' (384 dimensions, better quality)
+     * - 'Xenova/bge-base-en-v1.5' (768 dimensions, best quality)
+     * @default 'Xenova/all-MiniLM-L6-v2'
+     */
+    model?: string;
+
+    /**
+     * Optional. Maximum number of tokens that can be sent to the embedding model.
+     * @remarks
+     * This affects batching behavior in LocalDocumentIndex.
+     * Most small models support 512 tokens.
+     * @default 512
+     */
+    maxTokens?: number;
+
+    /**
+     * Optional. Device to run inference on.
+     * @remarks
+     * - 'auto': Automatically select the best available device
+     * - 'gpu': Use GPU (WebGPU in browser, CUDA in Node.js if available)
+     * - 'cpu': Use CPU (most compatible)
+     * - 'wasm': Use WebAssembly
+     * @default 'auto'
+     */
+    device?: 'auto' | 'gpu' | 'cpu' | 'wasm';
+
+    /**
+     * Optional. Data type for model weights.
+     * @remarks
+     * - 'fp32': Full precision (best quality, largest size)
+     * - 'fp16': Half precision (good quality, smaller)
+     * - 'q8': 8-bit quantization (good quality, smaller)
+     * - 'q4': 4-bit quantization (fastest, smallest, lower quality)
+     * @default 'fp32'
+     */
+    dtype?: 'fp32' | 'fp16' | 'q8' | 'q4';
+
+    /**
+     * Optional. Whether to normalize embeddings to unit length.
+     * @default true
+     */
+    normalize?: boolean;
+
+    /**
+     * Optional. Pooling strategy for token embeddings.
+     * @remarks
+     * - 'mean': Mean pooling (default, recommended)
+     * - 'cls': Use [CLS] token embedding
+     * @default 'mean'
+     */
+    pooling?: 'mean' | 'cls';
+
+    /**
+     * Optional. Callback for tracking model download/load progress.
+     */
+    progressCallback?: (progress: { status: string; progress?: number; file?: string }) => void;
+}
+
+/**
+ * An embeddings model using Transformers.js for local, offline inference.
+ * @remarks
+ * Requires @huggingface/transformers as a peer dependency.
+ * Use the static `create()` method to instantiate.
+ *
+ * @example
+ * ```typescript
+ * const embeddings = await TransformersEmbeddings.create({
+ *     model: 'Xenova/all-MiniLM-L6-v2'
+ * });
+ *
+ * const index = new LocalDocumentIndex({
+ *     folderPath: 'my-index',
+ *     embeddings: embeddings,
+ *     tokenizer: embeddings.getTokenizer()
+ * });
+ * ```
+ */
+export class TransformersEmbeddings implements EmbeddingsModel {
+    private readonly _extractor: FeatureExtractionPipeline;
+    private readonly _tokenizer: PreTrainedTokenizer;
+    private readonly _options: Required<Omit<TransformersEmbeddingsOptions, 'progressCallback'>> & Pick<TransformersEmbeddingsOptions, 'progressCallback'>;
+
+    public readonly maxTokens: number;
+
+    /**
+     * Private constructor - use TransformersEmbeddings.create() instead.
+     */
+    private constructor(
+        extractor: FeatureExtractionPipeline,
+        tokenizer: PreTrainedTokenizer,
+        options: Required<Omit<TransformersEmbeddingsOptions, 'progressCallback'>> & Pick<TransformersEmbeddingsOptions, 'progressCallback'>
+    ) {
+        this._extractor = extractor;
+        this._tokenizer = tokenizer;
+        this._options = options;
+        this.maxTokens = options.maxTokens;
+    }
+
+    /**
+     * Creates a new TransformersEmbeddings instance.
+     * @param options Configuration options.
+     * @returns Promise resolving to initialized TransformersEmbeddings instance.
+     * @throws Error if @huggingface/transformers is not installed.
+     */
+    public static async create(options?: TransformersEmbeddingsOptions): Promise<TransformersEmbeddings> {
+        // Dynamically import to allow optional dependency
+        let transformers: TransformersLibrary;
+
+        try {
+            transformers = await import('@huggingface/transformers');
+        } catch (e) {
+            throw new Error(
+                'TransformersEmbeddings requires @huggingface/transformers. ' +
+                'Install it with: npm install @huggingface/transformers'
+            );
+        }
+
+        const { pipeline } = transformers;
+
+        // Apply defaults
+        const opts = {
+            model: options?.model ?? DEFAULT_MODEL,
+            maxTokens: options?.maxTokens ?? 512,
+            device: options?.device ?? 'auto',
+            dtype: options?.dtype ?? 'fp32',
+            normalize: options?.normalize ?? true,
+            pooling: options?.pooling ?? 'mean',
+            progressCallback: options?.progressCallback
+        };
+
+        // Build pipeline options
+        const pipelineOptions: any = {
+            device: opts.device,
+            dtype: opts.dtype
+        };
+
+        if (opts.progressCallback) {
+            pipelineOptions.progress_callback = opts.progressCallback;
+        }
+
+        // Load the feature extraction pipeline
+        const extractor = await pipeline(
+            'feature-extraction',
+            opts.model,
+            pipelineOptions
+        );
+
+        // Load the tokenizer separately for use with TextSplitter
+        const tokenizer = extractor.tokenizer;
+
+        return new TransformersEmbeddings(extractor, tokenizer, opts);
+    }
+
+    /**
+     * Returns a tokenizer that uses the same tokenization as this embedding model.
+     * @remarks
+     * Use this tokenizer with LocalDocumentIndex to ensure text chunking
+     * aligns with the embedding model's token boundaries.
+     * @returns TransformersTokenizer instance.
+     */
+    public getTokenizer(): TransformersTokenizer {
+        return new TransformersTokenizer(this._tokenizer);
+    }
+
+    /**
+     * Creates embeddings for the given inputs.
+     * @param inputs Text inputs to create embeddings for.
+     * @returns EmbeddingsResponse with status and generated embeddings.
+     */
+    public async createEmbeddings(inputs: string | string[]): Promise<EmbeddingsResponse> {
+        try {
+            const inputArray = Array.isArray(inputs) ? inputs : [inputs];
+
+            // Process all inputs in a single batch
+            const output = await this._extractor(inputArray, {
+                pooling: this._options.pooling,
+                normalize: this._options.normalize
+            });
+
+            const [batchSize, embeddingDim] = output.dims;
+            const data = output.data as Float32Array;
+
+            // Slice the flat array into individual embeddings
+            const embeddings: number[][] = [];
+            for (let i = 0; i < batchSize; i++) {
+                const start = i * embeddingDim;
+                const end = start + embeddingDim;
+                embeddings.push(Array.from(data.slice(start, end)));
+            }
+
+            return {
+                status: 'success',
+                output: embeddings,
+                model: this._options.model
+            };
+        } catch (error: unknown) {
+            return {
+                status: 'error',
+                message: `Error generating embeddings: ${(error as Error).message}`
+            };
+        }
+    }
+
+    /**
+     * Returns the model name being used.
+     */
+    public get model(): string {
+        return this._options.model;
+    }
+}

package/src/TransformersTokenizer.spec.ts

@@ -0,0 +1,143 @@
+import { strict as assert } from 'node:assert';
+import { describe, it } from 'mocha';
+import { TransformersTokenizer } from './TransformersTokenizer';
+
+describe('TransformersTokenizer', () => {
+    // Create a mock tokenizer that mimics Transformers.js behavior
+    function createMockTokenizer() {
+        const vocab: Map<string, number> = new Map([
+            ['hello', 101],
+            ['world', 102],
+            ['test', 103],
+            ['[CLS]', 1],
+            ['[SEP]', 2]
+        ]);
+        const reverseVocab: Map<number, string> = new Map();
+        vocab.forEach((v, k) => reverseVocab.set(v, k));
+
+        return {
+            // Mimics the callable tokenizer behavior
+            __call__: (text: string) => {
+                const words = text.toLowerCase().split(/\s+/).filter(w => w);
+                const ids = words.map(w => vocab.get(w) ?? 100);
+                return {
+                    input_ids: {
+                        data: BigInt64Array.from(ids.map(id => BigInt(id)))
+                    }
+                };
+            },
+            decode: (tokens: number[], options?: { skip_special_tokens?: boolean }) => {
+                const words = tokens
+                    .filter(t => !options?.skip_special_tokens || (t !== 1 && t !== 2))
+                    .map(t => reverseVocab.get(t) ?? '[UNK]');
+                return words.join(' ');
+            }
+        };
+    }
+
+    it('encodes text to token array using callable tokenizer', () => {
+        const mockTokenizer = createMockTokenizer();
+        // Make it callable
+        const callableTokenizer = Object.assign(
+            (text: string) => mockTokenizer.__call__(text),
+            { decode: mockTokenizer.decode }
+        ) as any;
+
+        const tokenizer = new TransformersTokenizer(callableTokenizer);
+        const tokens = tokenizer.encode('hello world');
+
+        assert.ok(Array.isArray(tokens), 'encode should return an array');
+        assert.equal(tokens.length, 2, 'should have 2 tokens');
+        assert.deepEqual(tokens, [101, 102], 'tokens should match expected values');
+    });
+
+    it('handles BigInt64Array conversion correctly', () => {
+        const mockTokenizer = {
+            __call__: () => ({
+                input_ids: {
+                    data: BigInt64Array.from([BigInt(1), BigInt(2), BigInt(3)])
+                }
+            }),
+            decode: () => 'decoded'
+        };
+        const callableTokenizer = Object.assign(
+            () => mockTokenizer.__call__(),
+            { decode: mockTokenizer.decode }
+        ) as any;
+
+        const tokenizer = new TransformersTokenizer(callableTokenizer);
+        const tokens = tokenizer.encode('any text');
+
+        assert.deepEqual(tokens, [1, 2, 3], 'should convert BigInt to number');
+        tokens.forEach(t => {
+            assert.equal(typeof t, 'number', 'each token should be a number');
+        });
+    });
+
+    it('decodes tokens back to text', () => {
+        const mockTokenizer = {
+            __call__: () => ({ input_ids: { data: BigInt64Array.from([]) } }),
+            decode: (tokens: number[], opts?: { skip_special_tokens?: boolean }) => {
+                if (opts?.skip_special_tokens) {
+                    return 'hello world';
+                }
+                return '[CLS] hello world [SEP]';
+            }
+        };
+        const callableTokenizer = Object.assign(
+            () => mockTokenizer.__call__(),
+            { decode: mockTokenizer.decode }
+        ) as any;
+
+        const tokenizer = new TransformersTokenizer(callableTokenizer);
+        const text = tokenizer.decode([1, 101, 102, 2]);
+
+        assert.equal(text, 'hello world', 'should decode with skip_special_tokens=true');
+    });
+
+    it('handles empty input', () => {
+        const mockTokenizer = {
+            __call__: () => ({
+                input_ids: { data: BigInt64Array.from([]) }
+            }),
+            decode: () => ''
+        };
+        const callableTokenizer = Object.assign(
+            () => mockTokenizer.__call__(),
+            { decode: mockTokenizer.decode }
+        ) as any;
+
+        const tokenizer = new TransformersTokenizer(callableTokenizer);
+
+        const tokens = tokenizer.encode('');
+        assert.deepEqual(tokens, [], 'empty input should return empty array');
+
+        const text = tokenizer.decode([]);
+        assert.equal(text, '', 'empty tokens should return empty string');
+    });
+
+    it('returns consistent results for same input', () => {
+        let callCount = 0;
+        const mockTokenizer = {
+            __call__: () => {
+                callCount++;
+                return {
+                    input_ids: { data: BigInt64Array.from([BigInt(101), BigInt(102)]) }
+                };
+            },
+            decode: () => 'hello world'
+        };
+        const callableTokenizer = Object.assign(
+            () => mockTokenizer.__call__(),
+            { decode: mockTokenizer.decode }
+        ) as any;
+
+        const tokenizer = new TransformersTokenizer(callableTokenizer);
+
+        const tokens1 = tokenizer.encode('hello world');
+        const tokens2 = tokenizer.encode('hello world');
+
+        assert.deepEqual(tokens1, tokens2, 'encode should be deterministic');
+        assert.equal(callCount, 2, 'should call underlying tokenizer each time');
+    });
+});

package/src/TransformersTokenizer.ts

@@ -0,0 +1,45 @@
+import { PreTrainedTokenizer } from "@huggingface/transformers";
+import { Tokenizer } from "./types";
+
+/**
+ * A tokenizer wrapper for Transformers.js models.
+ * @remarks
+ * This tokenizer uses the same tokenizer as the embedding model,
+ * ensuring consistency between text splitting and embedding generation.
+ *
+ * Obtain an instance via TransformersEmbeddings.getTokenizer().
+ */
+export class TransformersTokenizer implements Tokenizer {
+    private readonly _tokenizer: PreTrainedTokenizer;
+
+    /**
+     * Creates a new TransformersTokenizer.
+     * @param tokenizer The underlying Transformers.js tokenizer.
+     * @remarks
+     * Typically created via TransformersEmbeddings.getTokenizer().
+     */
+    public constructor(tokenizer: PreTrainedTokenizer) {
+        this._tokenizer = tokenizer;
+    }
+
+    /**
+     * Encodes text into token IDs.
+     * @param text The text to encode.
+     * @returns Array of token IDs.
+     */
+    public encode(text: string): number[] {
+        const encoded = this._tokenizer(text);
+        // Transformers.js returns an object with input_ids as BigInt64Array or similar
+        const inputIds = encoded.input_ids?.data ?? encoded.input_ids ?? encoded;
+        return Array.from(inputIds).map((id: any) => Number(id));
+    }
+
+    /**
+     * Decodes token IDs back into text.
+     * @param tokens Array of token IDs.
+     * @returns Decoded text string.
+     */
+    public decode(tokens: number[]): string {
+        return this._tokenizer.decode(tokens, { skip_special_tokens: true });
+    }
+}