vectra 0.12.3 → 0.14.0

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 });
+    }
+}