David Snelling
f8c45f2d8d
Open source vector database with HNSW indexing, graph relationships, and metadata facets. Features CLI with professional augmentation registry integration for discovering extensions and capabilities.
409 lines
No EOL
16 KiB
JavaScript
409 lines
No EOL
16 KiB
JavaScript
/**
|
|
* Embedding functions for converting data to vectors using Transformers.js
|
|
* Complete rewrite to eliminate TensorFlow.js and use ONNX-based models
|
|
*/
|
|
import { isBrowser } from './environment.js';
|
|
// @ts-ignore - Transformers.js is now the primary embedding library
|
|
import { pipeline, env } from '@huggingface/transformers';
|
|
/**
|
|
* Detect the best available GPU device for the current environment
|
|
*/
|
|
export async function detectBestDevice() {
|
|
// Browser environment - check for WebGPU support
|
|
if (isBrowser()) {
|
|
if (typeof navigator !== 'undefined' && 'gpu' in navigator) {
|
|
try {
|
|
const adapter = await navigator.gpu?.requestAdapter();
|
|
if (adapter) {
|
|
return 'webgpu';
|
|
}
|
|
}
|
|
catch (error) {
|
|
// WebGPU not available or failed to initialize
|
|
}
|
|
}
|
|
return 'cpu';
|
|
}
|
|
// Node.js environment - check for CUDA support
|
|
try {
|
|
// Check if ONNX Runtime GPU packages are available
|
|
// This is a simple heuristic - in production you might want more sophisticated detection
|
|
const hasGpu = process.env.CUDA_VISIBLE_DEVICES !== undefined ||
|
|
process.env.ONNXRUNTIME_GPU_ENABLED === 'true';
|
|
return hasGpu ? 'cuda' : 'cpu';
|
|
}
|
|
catch (error) {
|
|
return 'cpu';
|
|
}
|
|
}
|
|
/**
|
|
* Resolve device string to actual device configuration
|
|
*/
|
|
export async function resolveDevice(device = 'auto') {
|
|
if (device === 'auto') {
|
|
return await detectBestDevice();
|
|
}
|
|
// Map 'gpu' to appropriate GPU type for current environment
|
|
if (device === 'gpu') {
|
|
const detected = await detectBestDevice();
|
|
return detected === 'cpu' ? 'cpu' : detected;
|
|
}
|
|
return device;
|
|
}
|
|
export class TransformerEmbedding {
|
|
/**
|
|
* Create a new TransformerEmbedding instance
|
|
*/
|
|
constructor(options = {}) {
|
|
this.extractor = null;
|
|
this.initialized = false;
|
|
this.verbose = true;
|
|
this.verbose = options.verbose !== undefined ? options.verbose : true;
|
|
// PRODUCTION-READY MODEL CONFIGURATION
|
|
// Priority order: explicit option > environment variable > smart default
|
|
let localFilesOnly;
|
|
if (options.localFilesOnly !== undefined) {
|
|
// 1. Explicit option takes highest priority
|
|
localFilesOnly = options.localFilesOnly;
|
|
}
|
|
else if (process.env.BRAINY_ALLOW_REMOTE_MODELS !== undefined) {
|
|
// 2. Environment variable override
|
|
localFilesOnly = process.env.BRAINY_ALLOW_REMOTE_MODELS !== 'true';
|
|
}
|
|
else if (process.env.NODE_ENV === 'development') {
|
|
// 3. Development mode allows remote models
|
|
localFilesOnly = false;
|
|
}
|
|
else if (isBrowser()) {
|
|
// 4. Browser defaults to allowing remote models
|
|
localFilesOnly = false;
|
|
}
|
|
else {
|
|
// 5. Node.js production: try local first, but allow remote as fallback
|
|
// This is the NEW production-friendly default
|
|
localFilesOnly = false;
|
|
}
|
|
this.options = {
|
|
model: options.model || 'Xenova/all-MiniLM-L6-v2',
|
|
verbose: this.verbose,
|
|
cacheDir: options.cacheDir || './models',
|
|
localFilesOnly: localFilesOnly,
|
|
dtype: options.dtype || 'fp32',
|
|
device: options.device || 'auto'
|
|
};
|
|
if (this.verbose) {
|
|
this.logger('log', `Embedding config: localFilesOnly=${localFilesOnly}, model=${this.options.model}, cacheDir=${this.options.cacheDir}`);
|
|
}
|
|
// Configure transformers.js environment
|
|
if (!isBrowser()) {
|
|
// Set cache directory for Node.js
|
|
env.cacheDir = this.options.cacheDir;
|
|
// Prioritize local models for offline operation
|
|
env.allowRemoteModels = !this.options.localFilesOnly;
|
|
env.allowLocalModels = true;
|
|
}
|
|
else {
|
|
// Browser configuration
|
|
// Allow both local and remote models, but prefer local if available
|
|
env.allowLocalModels = true;
|
|
env.allowRemoteModels = true;
|
|
// Force the configuration to ensure it's applied
|
|
if (this.verbose) {
|
|
this.logger('log', `Browser env config - allowLocalModels: ${env.allowLocalModels}, allowRemoteModels: ${env.allowRemoteModels}, localFilesOnly: ${this.options.localFilesOnly}`);
|
|
}
|
|
}
|
|
}
|
|
/**
|
|
* Get the default cache directory for models
|
|
*/
|
|
async getDefaultCacheDir() {
|
|
if (isBrowser()) {
|
|
return './models'; // Browser default
|
|
}
|
|
// Check for bundled models in the package
|
|
const possiblePaths = [
|
|
// In the installed package
|
|
'./node_modules/@soulcraft/brainy/models',
|
|
// In development/source
|
|
'./models',
|
|
'./dist/../models',
|
|
// Alternative locations
|
|
'../models',
|
|
'../../models'
|
|
];
|
|
// Check if we're in Node.js and try to find the bundled models
|
|
if (typeof process !== 'undefined' && process.versions?.node) {
|
|
try {
|
|
// Use dynamic import instead of require for ES modules compatibility
|
|
const { createRequire } = await import('module');
|
|
const require = createRequire(import.meta.url);
|
|
const path = require('path');
|
|
const fs = require('fs');
|
|
// Try to resolve the package location
|
|
try {
|
|
const brainyPackagePath = require.resolve('@soulcraft/brainy/package.json');
|
|
const brainyPackageDir = path.dirname(brainyPackagePath);
|
|
const bundledModelsPath = path.join(brainyPackageDir, 'models');
|
|
if (fs.existsSync(bundledModelsPath)) {
|
|
this.logger('log', `Using bundled models from package: ${bundledModelsPath}`);
|
|
return bundledModelsPath;
|
|
}
|
|
}
|
|
catch (e) {
|
|
// Not installed as package, continue
|
|
}
|
|
// Try relative paths from current location
|
|
for (const relativePath of possiblePaths) {
|
|
const fullPath = path.resolve(relativePath);
|
|
if (fs.existsSync(fullPath)) {
|
|
this.logger('log', `Using bundled models from: ${fullPath}`);
|
|
return fullPath;
|
|
}
|
|
}
|
|
}
|
|
catch (error) {
|
|
// Silently fall back to default path if module detection fails
|
|
}
|
|
}
|
|
// Fallback to default cache directory
|
|
return './models';
|
|
}
|
|
/**
|
|
* Check if we're running in a test environment
|
|
*/
|
|
isTestEnvironment() {
|
|
// Always use real implementation - no more mocking
|
|
return false;
|
|
}
|
|
/**
|
|
* Log message only if verbose mode is enabled
|
|
*/
|
|
logger(level, message, ...args) {
|
|
if (level === 'error' || this.verbose) {
|
|
console[level](`[TransformerEmbedding] ${message}`, ...args);
|
|
}
|
|
}
|
|
/**
|
|
* Initialize the embedding model
|
|
*/
|
|
async init() {
|
|
if (this.initialized) {
|
|
return;
|
|
}
|
|
// Always use real implementation - no mocking
|
|
try {
|
|
// Resolve device configuration and cache directory
|
|
const device = await resolveDevice(this.options.device);
|
|
const cacheDir = this.options.cacheDir === './models'
|
|
? await this.getDefaultCacheDir()
|
|
: this.options.cacheDir;
|
|
this.logger('log', `Loading Transformer model: ${this.options.model} on device: ${device}`);
|
|
const startTime = Date.now();
|
|
// Load the feature extraction pipeline with GPU support
|
|
const pipelineOptions = {
|
|
cache_dir: cacheDir,
|
|
local_files_only: isBrowser() ? false : this.options.localFilesOnly,
|
|
dtype: this.options.dtype
|
|
};
|
|
// Add device configuration for GPU acceleration
|
|
if (device !== 'cpu') {
|
|
pipelineOptions.device = device;
|
|
this.logger('log', `🚀 GPU acceleration enabled: ${device}`);
|
|
}
|
|
if (this.verbose) {
|
|
this.logger('log', `Pipeline options: ${JSON.stringify(pipelineOptions)}`);
|
|
}
|
|
try {
|
|
this.extractor = await pipeline('feature-extraction', this.options.model, pipelineOptions);
|
|
}
|
|
catch (gpuError) {
|
|
// Fallback to CPU if GPU initialization fails
|
|
if (device !== 'cpu') {
|
|
this.logger('warn', `GPU initialization failed, falling back to CPU: ${gpuError?.message || gpuError}`);
|
|
const cpuOptions = { ...pipelineOptions };
|
|
delete cpuOptions.device;
|
|
this.extractor = await pipeline('feature-extraction', this.options.model, cpuOptions);
|
|
}
|
|
else {
|
|
// PRODUCTION-READY ERROR HANDLING
|
|
// If local_files_only is true and models are missing, try enabling remote downloads
|
|
if (pipelineOptions.local_files_only && gpuError?.message?.includes('local_files_only')) {
|
|
this.logger('warn', 'Local models not found, attempting remote download as fallback...');
|
|
try {
|
|
const remoteOptions = { ...pipelineOptions, local_files_only: false };
|
|
this.extractor = await pipeline('feature-extraction', this.options.model, remoteOptions);
|
|
this.logger('log', '✅ Successfully downloaded and loaded model from remote');
|
|
// Update the configuration to reflect what actually worked
|
|
this.options.localFilesOnly = false;
|
|
}
|
|
catch (remoteError) {
|
|
// Both local and remote failed - throw comprehensive error
|
|
const errorMsg = `Failed to load embedding model "${this.options.model}". ` +
|
|
`Local models not found and remote download failed. ` +
|
|
`To fix: 1) Set BRAINY_ALLOW_REMOTE_MODELS=true, ` +
|
|
`2) Run "npm run download-models", or ` +
|
|
`3) Use a custom embedding function.`;
|
|
throw new Error(errorMsg);
|
|
}
|
|
}
|
|
else {
|
|
throw gpuError;
|
|
}
|
|
}
|
|
}
|
|
const loadTime = Date.now() - startTime;
|
|
this.logger('log', `✅ Model loaded successfully in ${loadTime}ms`);
|
|
this.initialized = true;
|
|
}
|
|
catch (error) {
|
|
this.logger('error', 'Failed to initialize Transformer embedding model:', error);
|
|
throw new Error(`Transformer embedding initialization failed: ${error}`);
|
|
}
|
|
}
|
|
/**
|
|
* Generate embeddings for text data
|
|
*/
|
|
async embed(data) {
|
|
if (!this.initialized) {
|
|
await this.init();
|
|
}
|
|
try {
|
|
// Handle different input types
|
|
let textToEmbed;
|
|
if (typeof data === 'string') {
|
|
// Handle empty string case
|
|
if (data.trim() === '') {
|
|
// Return a zero vector of 384 dimensions (all-MiniLM-L6-v2 standard)
|
|
return new Array(384).fill(0);
|
|
}
|
|
textToEmbed = [data];
|
|
}
|
|
else if (Array.isArray(data) && data.every((item) => typeof item === 'string')) {
|
|
// Handle empty array or array with empty strings
|
|
if (data.length === 0 || data.every((item) => item.trim() === '')) {
|
|
return new Array(384).fill(0);
|
|
}
|
|
// Filter out empty strings
|
|
textToEmbed = data.filter((item) => item.trim() !== '');
|
|
if (textToEmbed.length === 0) {
|
|
return new Array(384).fill(0);
|
|
}
|
|
}
|
|
else {
|
|
throw new Error('TransformerEmbedding only supports string or string[] data');
|
|
}
|
|
// Ensure the extractor is available
|
|
if (!this.extractor) {
|
|
throw new Error('Transformer embedding model is not available');
|
|
}
|
|
// Generate embeddings with mean pooling and normalization
|
|
const result = await this.extractor(textToEmbed, {
|
|
pooling: 'mean',
|
|
normalize: true
|
|
});
|
|
// Extract the embedding data
|
|
let embedding;
|
|
if (textToEmbed.length === 1) {
|
|
// Single text input - return first embedding
|
|
embedding = Array.from(result.data.slice(0, 384));
|
|
}
|
|
else {
|
|
// Multiple texts - return first embedding (maintain compatibility)
|
|
embedding = Array.from(result.data.slice(0, 384));
|
|
}
|
|
// Validate embedding dimensions
|
|
if (embedding.length !== 384) {
|
|
this.logger('warn', `Unexpected embedding dimension: ${embedding.length}, expected 384`);
|
|
// Pad or truncate to 384 dimensions
|
|
if (embedding.length < 384) {
|
|
embedding = [...embedding, ...new Array(384 - embedding.length).fill(0)];
|
|
}
|
|
else {
|
|
embedding = embedding.slice(0, 384);
|
|
}
|
|
}
|
|
return embedding;
|
|
}
|
|
catch (error) {
|
|
this.logger('error', 'Error generating embeddings:', error);
|
|
throw new Error(`Failed to generate embeddings: ${error}`);
|
|
}
|
|
}
|
|
/**
|
|
* Dispose of the model and free resources
|
|
*/
|
|
async dispose() {
|
|
if (this.extractor && typeof this.extractor.dispose === 'function') {
|
|
await this.extractor.dispose();
|
|
}
|
|
this.extractor = null;
|
|
this.initialized = false;
|
|
}
|
|
/**
|
|
* Get the dimension of embeddings produced by this model
|
|
*/
|
|
getDimension() {
|
|
return 384;
|
|
}
|
|
/**
|
|
* Check if the model is initialized
|
|
*/
|
|
isInitialized() {
|
|
return this.initialized;
|
|
}
|
|
}
|
|
// Legacy alias for backward compatibility
|
|
export const UniversalSentenceEncoder = TransformerEmbedding;
|
|
/**
|
|
* Create a new embedding model instance
|
|
*/
|
|
export function createEmbeddingModel(options) {
|
|
return new TransformerEmbedding(options);
|
|
}
|
|
/**
|
|
* Default embedding function using the lightweight transformer model
|
|
*/
|
|
export const defaultEmbeddingFunction = async (data) => {
|
|
const embedder = new TransformerEmbedding({ verbose: false });
|
|
return await embedder.embed(data);
|
|
};
|
|
/**
|
|
* Create an embedding function with custom options
|
|
*/
|
|
export function createEmbeddingFunction(options = {}) {
|
|
const embedder = new TransformerEmbedding(options);
|
|
return async (data) => {
|
|
return await embedder.embed(data);
|
|
};
|
|
}
|
|
/**
|
|
* Batch embedding function for processing multiple texts efficiently
|
|
*/
|
|
export async function batchEmbed(texts, options = {}) {
|
|
const embedder = new TransformerEmbedding(options);
|
|
await embedder.init();
|
|
const embeddings = [];
|
|
// Process in batches for memory efficiency
|
|
const batchSize = 32;
|
|
for (let i = 0; i < texts.length; i += batchSize) {
|
|
const batch = texts.slice(i, i + batchSize);
|
|
for (const text of batch) {
|
|
const embedding = await embedder.embed(text);
|
|
embeddings.push(embedding);
|
|
}
|
|
}
|
|
await embedder.dispose();
|
|
return embeddings;
|
|
}
|
|
/**
|
|
* Embedding functions for specific model types
|
|
*/
|
|
export const embeddingFunctions = {
|
|
/** Default lightweight model (all-MiniLM-L6-v2, 384 dimensions) */
|
|
default: defaultEmbeddingFunction,
|
|
/** Create custom embedding function */
|
|
create: createEmbeddingFunction,
|
|
/** Batch processing */
|
|
batch: batchEmbed
|
|
};
|
|
//# sourceMappingURL=embedding.js.map
|