From bba9a0c219c308f82821e46c44c49ef560ca2e39 Mon Sep 17 00:00:00 2001 From: David Snelling Date: Fri, 27 Jun 2025 14:06:59 -0700 Subject: [PATCH] feat(src/brainyData, src/utils): enhance embedding efficiency with batch processing and initialize safeguards - Added batch embedding support with `defaultBatchEmbeddingFunction`, leveraging shared model instances for optimized performance. - Integrated `isInitializing` flag to prevent recursive initialization and ensure smooth concurrent operation handling during `BrainyData` initialization. - Pre-loaded Universal Sentence Encoder in `BrainyData` to prevent delays during embedding. - Introduced fallback mechanisms in embedding initialization for better error resiliency and model reusability. - Updated `addBatch` with support for batchSize and refactored text/vector processing logic for clearer separation and memory management. - Improved GPU and CPU backend selection in Universal Sentence Encoder for compatibility across environments. - Enhanced memory management by cleaning tensors after embedding operations. - Updated README with instructions for batch embedding, threading updates, and GPU/CPU optimizations. --- README.md | 57 ++--- demo/index.html | 152 +++++++------ package.json | 1 + src/brainyData.ts | 155 +++++++++++++- src/global.d.ts | 15 +- src/hnsw/hnswIndex.ts | 52 +---- src/utils/distance.ts | 310 ++++++++++++++------------- src/utils/embedding.ts | 452 +++++++++++++++++++++++---------------- src/utils/environment.ts | 5 +- src/utils/index.ts | 1 + src/utils/workerUtils.ts | 162 ++------------ 11 files changed, 729 insertions(+), 633 deletions(-) diff --git a/README.md b/README.md index c40e48bd..6d7bcfa5 100644 --- a/README.md +++ b/README.md @@ -508,7 +508,7 @@ const id = await db.add(textOrVector, { // other metadata... }) -// Add multiple nouns in parallel (with multithreading) +// Add multiple nouns in parallel (with multithreading and batch embedding) const ids = await db.addBatch([ { vectorOrData: "First item to add", @@ -521,7 +521,8 @@ const ids = await db.addBatch([ // More items... ], { forceEmbed: false, - concurrency: 4 // Control the level of parallelism (default: 4) + concurrency: 4, // Control the level of parallelism (default: 4) + batchSize: 50 // Control the number of items to process in a single batch (default: 50) }) // Retrieve a noun @@ -591,9 +592,7 @@ await db.init() // Or use the threaded embedding function for better performance const threadedDb = new BrainyData({ - embeddingFunction: createThreadedEmbeddingFunction({ - fallbackToMain: true // Fall back to main thread if threading fails - }) + embeddingFunction: createThreadedEmbeddingFunction() }) await threadedDb.init() @@ -602,21 +601,24 @@ const vector = await db.embed("Some text to convert to a vector") ``` The threaded embedding function runs in a separate thread (Web Worker in browsers, Worker Thread in Node.js) to improve -performance, especially for CPU-intensive embedding operations. It automatically falls back to the main thread if -threading is not available in the current environment. +performance, especially for embedding operations. It uses GPU acceleration when available (via WebGL in browsers) and +falls back to CPU processing for compatibility. Universal Sentence Encoder is always used for embeddings. The +implementation includes worker reuse and model caching for optimal performance. ### Performance Tuning -Brainy includes comprehensive performance optimizations that work across all environments (browser, CLI, Node.js, container, server): +Brainy includes comprehensive performance optimizations that work across all environments (browser, CLI, Node.js, +container, server): -#### GPU Acceleration +#### GPU and CPU Optimization -Brainy now leverages GPU acceleration via WebGL for compute-intensive operations: +Brainy uses GPU and CPU optimization for compute-intensive operations: -1. **GPU-Accelerated Embeddings**: Generate text embeddings using TensorFlow.js with WebGL backend -2. **GPU-Accelerated Distance Calculations**: Perform vector similarity calculations on the GPU for faster search -3. **Automatic Fallback**: Falls back to CPU processing when GPU is not available or fails to initialize -4. **Cross-Environment Support**: Works in browsers (via WebGL) and Node.js environments +1. **GPU-Accelerated Embeddings**: Generate text embeddings using TensorFlow.js with WebGL backend when available +2. **Automatic Fallback**: Falls back to CPU backend when GPU is not available +3. **Optimized Distance Calculations**: Perform vector similarity calculations with optimized algorithms +4. **Cross-Environment Support**: Works consistently across browsers and Node.js environments +5. **Memory Management**: Properly disposes of tensors to prevent memory leaks #### Multithreading Support @@ -625,7 +627,10 @@ Brainy includes comprehensive multithreading support to improve performance acro 1. **Parallel Batch Processing**: Add multiple items concurrently with controlled parallelism 2. **Multithreaded Vector Search**: Perform distance calculations in parallel for faster search operations 3. **Threaded Embedding Generation**: Generate embeddings in separate threads to avoid blocking the main thread -4. **Automatic Environment Detection**: Adapts to browser (Web Workers) and Node.js (Worker Threads) environments +4. **Worker Reuse**: Maintains a pool of workers to avoid the overhead of creating and terminating workers +5. **Model Caching**: Initializes the embedding model once per worker and reuses it for multiple operations +6. **Batch Embedding**: Processes multiple items in a single embedding operation for better performance +7. **Automatic Environment Detection**: Adapts to browser (Web Workers) and Node.js (Worker Threads) environments ```typescript import { BrainyData, euclideanDistance } from '@soulcraft/brainy' @@ -645,8 +650,6 @@ const db = new BrainyData({ // Performance optimization options performance: { useParallelization: true, // Enable multithreaded search operations - useGPUAcceleration: true, // Enable GPU acceleration for compute-intensive operations - fallbackToCPU: true, // Fall back to CPU processing if GPU acceleration fails }, // Noun and Verb type validation @@ -728,12 +731,9 @@ Brainy provides several distance functions for vector similarity calculations: - `manhattanDistance`: Measures the sum of absolute differences between vector components - `dotProductDistance`: Measures the negative dot product between vectors -All distance functions have GPU-accelerated versions that are automatically used when: -1. GPU acceleration is enabled in the configuration -2. The operation involves a large number of vectors -3. WebGL is available in the environment - -The GPU-accelerated distance calculations provide significant performance improvements for large datasets and high-dimensional vectors. +All distance functions are optimized for performance and automatically use the most efficient implementation based on +the dataset size and available resources. For large datasets and high-dimensional vectors, Brainy uses batch processing +and multithreading when available to improve performance. ## Backup and Restore @@ -815,6 +815,9 @@ brainy import-sparse --input sparse-data.json Brainy uses the following embedding approach: - TensorFlow Universal Sentence Encoder (high-quality text embeddings) +- GPU acceleration when available (via WebGL in browsers) +- Batch embedding for processing multiple items efficiently +- Worker reuse and model caching for optimal performance - Custom embedding functions can be plugged in for specialized domains ## Extensions @@ -1240,7 +1243,7 @@ const id = await db.addToBoth('Deep learning is a subset of machine learning', { tags: ['deep learning', 'neural networks'] }) -// Clean up when done +// Clean up when done (this also cleans up worker pools) await db.shutDown() ``` @@ -1298,9 +1301,9 @@ Brainy follows a specific code style to maintain consistency throughout the code The README badges are automatically updated during the build process: 1. **npm Version Badge**: The npm version badge is automatically updated to match the version in package.json when: - - Running `npm run build` (via the prebuild script) - - Running `npm version` commands (patch, minor, major) - - Manually running `node scripts/generate-version.js` + - Running `npm run build` (via the prebuild script) + - Running `npm version` commands (patch, minor, major) + - Manually running `node scripts/generate-version.js` This ensures that the badge always reflects the current version in package.json, even before publishing to npm. diff --git a/demo/index.html b/demo/index.html index 7f8293cd..774e6f6d 100644 --- a/demo/index.html +++ b/demo/index.html @@ -1448,7 +1448,10 @@ await db.init() await database.init() - log('db-status', 'Database initialized successfully!\nStorage: OPFS\nAugmentations: Embedding, Search') + // Clear the database after initialization + await database.clear() + + log('db-status', 'Database initialized and cleared successfully!\nStorage: OPFS\nAugmentations: Embedding, Search') updateDatabaseVisualization() // Populate noun dropdowns after initialization @@ -1475,40 +1478,40 @@ await db.init() type: 'person', noun: 'person', name: 'Alex Johnson', - content: 'Tech enthusiast and software developer. Love hiking and photography. Working on AI projects in my spare time.', - tags: ['tech', 'developer', 'photography', 'hiking'] + content: 'Tech enthusiast and senior software developer with 10+ years of experience in AI and machine learning. I love hiking through national parks on weekends and capturing landscape photography. Currently working on several open-source AI projects that focus on natural language processing and computer vision applications for everyday use cases.', + tags: ['tech', 'developer', 'photography', 'hiking', 'AI', 'machine learning'] }, { id: 'user2', type: 'person', noun: 'person', name: 'Sophia Chen', - content: 'Digital marketing specialist with a passion for data analytics. Foodie and travel blogger on weekends.', - tags: ['marketing', 'analytics', 'food', 'travel'] + content: 'Digital marketing specialist with a passion for data analytics and consumer behavior research. I run a popular food and travel blog documenting culinary experiences across continents. My background in statistics helps me analyze food trends and cultural patterns in different regions. Recently completed a six-month journey through Southeast Asia exploring local cuisines.', + tags: ['marketing', 'analytics', 'food', 'travel', 'statistics', 'blogging'] }, { id: 'user3', type: 'person', noun: 'person', name: 'Marcus Williams', - content: 'Graphic designer and illustrator. Creating visual stories that connect people. Coffee addict.', - tags: ['design', 'illustration', 'art', 'coffee'] + content: 'Award-winning graphic designer and illustrator specializing in brand identity and visual storytelling. I believe in the power of minimalist design to communicate complex ideas effectively. My work has been featured in several international design publications and exhibitions. When not designing, I\'m exploring local coffee shops and perfecting my home brewing techniques.', + tags: ['design', 'illustration', 'art', 'coffee', 'branding', 'minimalism'] }, { id: 'user4', type: 'person', noun: 'person', name: 'Priya Patel', - content: 'Environmental scientist studying climate change impacts. Advocate for sustainable living and renewable energy.', - tags: ['science', 'environment', 'sustainability', 'climate'] + content: 'Environmental scientist with a PhD in climate science, currently researching the long-term impacts of climate change on urban ecosystems. I advocate for sustainable living practices and renewable energy adoption through community workshops and educational programs. My research has been published in several peer-reviewed journals, and I consult with local governments on sustainability initiatives.', + tags: ['science', 'environment', 'sustainability', 'climate', 'research', 'education'] }, { id: 'user5', type: 'person', noun: 'person', name: 'Jordan Taylor', - content: 'Fitness coach and nutrition expert. Helping people achieve their health goals through balanced lifestyle.', - tags: ['fitness', 'nutrition', 'health', 'wellness'] + content: 'Certified fitness coach and nutrition expert with specializations in functional training and plant-based nutrition. I develop personalized fitness programs that integrate physical training with mental wellness practices. My approach focuses on sustainable lifestyle changes rather than quick fixes. I\'ve helped hundreds of clients transform their health through balanced nutrition and consistent exercise routines.', + tags: ['fitness', 'nutrition', 'health', 'wellness', 'coaching', 'plant-based'] }, // Posts/Content @@ -1517,40 +1520,40 @@ await db.init() type: 'document', noun: 'content', title: 'The Future of AI in Everyday Applications', - content: 'Artificial intelligence is rapidly transforming how we interact with technology. From smart assistants to recommendation systems, AI is becoming an invisible part of our daily lives. What excites me most is how these technologies can be personalized to individual needs while respecting privacy.', - tags: ['AI', 'technology', 'future', 'personalization'] + content: 'Artificial intelligence is rapidly transforming how we interact with technology in ways that weren\'t possible just a few years ago. From smart assistants that understand context and nuance to recommendation systems that seem to know our preferences better than we do, AI is becoming an invisible yet powerful part of our daily lives. The most exciting developments are happening at the intersection of machine learning and edge computing, where AI models can run directly on our devices without sending sensitive data to the cloud. This preserves privacy while still delivering personalized experiences. Looking ahead, I believe we\'ll see AI becoming more explainable and transparent, addressing current concerns about "black box" algorithms. The challenge for developers will be creating systems that can be personalized to individual needs while maintaining ethical standards and avoiding algorithmic bias. The companies that succeed will be those that treat AI as an augmentation of human capabilities rather than a replacement.', + tags: ['AI', 'technology', 'future', 'personalization', 'machine learning', 'privacy', 'ethics'] }, { id: 'post2', type: 'document', noun: 'content', - title: 'My Southeast Asia Food Tour', - content: 'Just returned from an amazing culinary adventure across Thailand, Vietnam, and Malaysia. The street food scenes in Bangkok and Hanoi were incredible! Discovered so many new flavors and cooking techniques that I can\'t wait to try at home.', - tags: ['food', 'travel', 'asia', 'culinary'] + title: 'My Southeast Asia Food Tour: A Culinary Journey', + content: 'I\'ve just returned from an incredible three-month culinary adventure across Thailand, Vietnam, Malaysia, and Singapore that completely transformed my understanding of Asian cuisine. The vibrant street food scenes in Bangkok and Hanoi were particularly eye-opening - watching skilled vendors prepare dishes they\'ve perfected over decades, using techniques passed down through generations. In Bangkok\'s Chinatown, I discovered that the best pad thai comes from a 70-year-old woman who cooks only 50 portions each night, selling out within an hour. The complex interplay of sweet, sour, salty, and spicy flavors in authentic Thai cuisine is something that western adaptations rarely capture correctly. Moving to Vietnam, I was fascinated by the regional variations in pho - from the clearer, more delicate broths in the south to the richer, more aromatic versions in Hanoi. The Malaysian food scene in Penang revealed the beautiful fusion of Chinese, Indian, and Malay influences, creating dishes that tell the story of the country\'s multicultural history. I\'ve returned home with not just memories, but with cooking techniques, ingredient combinations, and flavor profiles that I\'m excited to incorporate into my own cooking. The experience has reinforced my belief that food is one of the most authentic ways to understand a culture\'s values and history.', + tags: ['food', 'travel', 'asia', 'culinary', 'culture', 'cooking', 'Thailand', 'Vietnam', 'Malaysia'] }, { id: 'post3', type: 'document', noun: 'content', - title: 'Minimalist Design Principles for Digital Interfaces', - content: 'Less is more when it comes to effective UI design. This article explores how minimalist principles can improve user experience, reduce cognitive load, and create more accessible digital products.', - tags: ['design', 'minimalism', 'UI', 'UX'] + title: 'Minimalist Design Principles for Digital Interfaces: Beyond Aesthetics', + content: 'The principle that "less is more" has never been more relevant than in today\'s overcrowded digital landscape. Minimalist design isn\'t just an aesthetic choice—it\'s a functional imperative for creating interfaces that users can navigate intuitively. Through my work with various clients, I\'ve observed that minimalist interfaces consistently outperform cluttered ones in key metrics like time-on-task, error rates, and user satisfaction. The core principles of minimalist design extend beyond simply removing elements. It involves a careful analysis of user needs, prioritizing functions based on frequency and importance, and creating visual hierarchies that guide attention naturally. Color should be used strategically, not decoratively—each color choice should serve a specific communicative purpose. Typography plays an equally crucial role; a well-chosen font hierarchy can eliminate the need for additional visual elements like dividers or boxes. White space (or negative space) is perhaps the most undervalued element in digital design. It\'s not just empty space—it\'s a powerful tool for creating focus, improving readability, and giving content room to breathe. When implementing minimalist principles, it\'s important to test with real users to ensure that simplification doesn\'t compromise usability. The goal is to reduce cognitive load while maintaining all necessary functionality, creating digital products that feel effortless to use while being fully featured. This approach not only improves user experience but also tends to create more accessible interfaces that work better for users with disabilities or those using assistive technologies.', + tags: ['design', 'minimalism', 'UI', 'UX', 'accessibility', 'typography', 'user experience', 'digital interfaces'] }, { id: 'post4', type: 'document', noun: 'content', title: 'Urban Gardening: Growing Food in Small Spaces', - content: 'You don\'t need a large yard to grow your own food. This guide shows how to create productive gardens in apartments, balconies, and small urban spaces using container gardening, vertical systems, and efficient planning.', - tags: ['gardening', 'sustainability', 'urban', 'food'] + content: 'The misconception that you need a large yard to grow your own food is preventing many urban dwellers from experiencing the satisfaction and benefits of home gardening. Through five years of experimentation in my 500-square-foot apartment, I\'ve developed systems for growing approximately 30% of my own produce year-round. Vertical gardening is the cornerstone of space-efficient food production. By utilizing wall space with hanging planters, tiered shelving, and trellis systems, even the smallest balcony can become surprisingly productive. I\'ve found that leafy greens like kale, spinach, and various lettuces offer the best return on investment in small spaces—they grow quickly, can be harvested continuously, and would otherwise be expensive to purchase organically. For those without outdoor space, windowsill herb gardens and microgreens under grow lights provide fresh flavors with minimal space requirements. The key to successful small-space gardening is understanding the specific light conditions of your space and selecting appropriate plants. Most vegetables need at least 6 hours of direct sunlight, but many herbs and leafy greens can thrive with less. Container selection is equally important—self-watering containers have revolutionized urban gardening by reducing maintenance and improving plant health. Soil quality becomes even more critical in container gardening; investing in high-quality potting mix and implementing a regular composting system (even a small worm bin under the sink) creates a sustainable cycle. Beyond the practical benefits of fresh, pesticide-free produce, urban gardening creates a connection to natural cycles that is often missing in city life. The psychological benefits of tending plants and watching them grow should not be underestimated, especially in high-stress urban environments.', + tags: ['gardening', 'sustainability', 'urban', 'food', 'containers', 'vertical gardening', 'organic', 'small spaces'] }, { id: 'post5', type: 'document', noun: 'content', - title: '30-Day Home Workout Challenge', - content: 'Transform your fitness with this comprehensive 30-day workout plan that requires no equipment. Each day builds on the previous one, gradually increasing intensity while focusing on different muscle groups.', - tags: ['fitness', 'workout', 'health', 'challenge'] + title: '30-Day Home Workout Challenge: Progressive Fitness Without Equipment', + content: 'After years of designing fitness programs, I\'ve found that the biggest obstacle for most people isn\'t motivation—it\'s complexity and accessibility. This 30-day challenge addresses both issues by requiring zero equipment while providing a structured progression that prevents plateaus and reduces injury risk. The program is built around functional movement patterns that strengthen the body for real-world activities, not just gym performance. The first week focuses on establishing proper form in fundamental movements like squats, lunges, planks, and push-ups, with modifications provided for all fitness levels. Rather than counting reps, many exercises use time-based intervals, allowing each person to work at their own pace while still being challenged. As the program progresses, intensity increases through three primary mechanisms: increasing time under tension, incorporating unilateral (single-limb) variations, and adding rhythmic changes that challenge coordination and cardiovascular fitness simultaneously. Recovery is programmed as deliberately as the workouts themselves, with specific mobility routines for rest days that address common problem areas like hip flexors and shoulders that become tight from sedentary work. The nutrition guidance accompanying the program emphasizes timing and composition rather than strict calorie counting—particularly the importance of protein intake for recovery and carbohydrate timing around workouts. Participants who\'ve completed this challenge report not just physical changes, but improvements in sleep quality, energy levels, and mental clarity. The program includes a tracking system that helps identify patterns between workout performance and variables like sleep, stress, and nutrition, teaching participants to understand their body\'s unique responses and needs.', + tags: ['fitness', 'workout', 'health', 'challenge', 'home exercise', 'bodyweight', 'functional fitness', 'progressive training'] }, // Comments @@ -1559,24 +1562,24 @@ await db.init() type: 'document', noun: 'content', title: 'Comment on AI post', - content: 'Great insights! I\'m particularly interested in how AI can be applied to healthcare. Have you explored that area?', - tags: ['comment', 'AI', 'healthcare'] + content: 'Your insights on the intersection of AI and privacy are particularly relevant given the recent controversies surrounding facial recognition technologies. I\'m currently researching how AI can be applied to healthcare diagnostics while maintaining patient confidentiality. Have you explored how federated learning might offer solutions in this domain? I\'d love to discuss this further as it seems aligned with your interest in edge computing applications.', + tags: ['comment', 'AI', 'healthcare', 'privacy', 'federated learning', 'edge computing'] }, { id: 'comment2', type: 'document', noun: 'content', title: 'Comment on food tour', - content: 'Your photos are amazing! What was your favorite dish from the entire trip?', - tags: ['comment', 'food', 'travel', 'photography'] + content: 'Your detailed description of regional Vietnamese pho variations brought back memories of my own travels there! The photos you shared of the Hanoi street food markets are absolutely stunning - the composition and lighting really capture the vibrant atmosphere. I\'m curious about your experience with Malaysian laksa - did you find significant differences between the Penang and Sarawak versions? I\'ve been trying to recreate authentic laksa at home but struggling to find the right balance of spices.', + tags: ['comment', 'food', 'travel', 'photography', 'Vietnam', 'Malaysia', 'cooking', 'recipes'] }, { id: 'comment3', type: 'document', noun: 'content', title: 'Comment on design article', - content: 'Minimalism is definitely underrated. I\'ve found that my conversion rates improved significantly after simplifying our landing page.', - tags: ['comment', 'design', 'conversion', 'business'] + content: 'Your analysis of minimalism as a functional imperative rather than just an aesthetic choice resonates strongly with my experience in e-commerce design. After implementing many of the principles you outlined, particularly regarding strategic color usage and improved typography hierarchy, our conversion rates improved by 23% while customer support inquiries decreased by almost 30%. The point about accessibility benefits is especially important and often overlooked in design discussions. Have you found any effective methods for convincing stakeholders who equate visual complexity with feature richness?', + tags: ['comment', 'design', 'conversion', 'business', 'accessibility', 'minimalism', 'typography', 'e-commerce'] }, // Groups/Communities @@ -1585,24 +1588,24 @@ await db.init() type: 'group', noun: 'group', name: 'Tech Innovators Network', - content: 'A community of developers, designers, and tech enthusiasts discussing emerging technologies and their applications.', - tags: ['technology', 'innovation', 'community', 'networking'] + content: 'A global community of developers, designers, and tech enthusiasts collaborating on emerging technologies and their practical applications. Our members range from startup founders to corporate innovators, all sharing insights on AI, blockchain, IoT, and other transformative technologies. We host monthly virtual meetups featuring expert speakers, maintain an active knowledge base of implementation case studies, and facilitate mentorship connections between experienced professionals and those entering the field.', + tags: ['technology', 'innovation', 'community', 'networking', 'AI', 'blockchain', 'IoT', 'collaboration'] }, { id: 'group2', type: 'group', noun: 'group', - name: 'Global Foodies', - content: 'Sharing culinary experiences, recipes, and food photography from around the world.', - tags: ['food', 'cooking', 'international', 'recipes'] + name: 'Global Foodies Collective', + content: 'An international community dedicated to exploring culinary traditions and innovations from around the world. Our members share authentic recipes, cooking techniques, and food photography that celebrates cultural diversity through cuisine. We organize themed cooking challenges, virtual cook-alongs with chefs from different countries, and discussions about sustainable food systems and the preservation of traditional cooking methods. The community also maintains a searchable database of member-tested recipes organized by region, dietary preferences, and difficulty level.', + tags: ['food', 'cooking', 'international', 'recipes', 'culture', 'sustainability', 'traditions', 'culinary'] }, { id: 'group3', type: 'group', noun: 'group', name: 'Sustainable Living Collective', - content: 'Exploring practical ways to reduce environmental impact and live more sustainably in modern society.', - tags: ['sustainability', 'environment', 'lifestyle', 'eco-friendly'] + content: 'A community focused on practical approaches to reducing environmental impact in everyday life. We share evidence-based strategies for sustainable living in urban, suburban, and rural contexts, with special emphasis on solutions that are economically accessible. Our resources include guides for zero-waste home management, energy efficiency improvements, ethical consumption, and small-space gardening. Members collaborate on neighborhood initiatives like community gardens, repair cafés, and local policy advocacy for environmental protection and green infrastructure.', + tags: ['sustainability', 'environment', 'lifestyle', 'eco-friendly', 'zero-waste', 'community', 'gardening', 'advocacy'] }, // Events @@ -1610,28 +1613,29 @@ await db.init() id: 'event1', type: 'event', noun: 'event', - name: 'Virtual Tech Meetup 2023', - content: 'Online gathering of tech professionals sharing insights on latest development frameworks and tools.', - tags: ['tech', 'virtual', 'networking', 'learning'] + name: 'Virtual Tech Meetup 2023: Responsible AI Development', + content: 'An online gathering of technology professionals sharing insights on ethical AI development frameworks and implementation tools. The event features keynote presentations on bias detection in machine learning models, panel discussions on regulatory compliance across different jurisdictions, and hands-on workshops demonstrating techniques for making AI systems more transparent and explainable. Participants will have opportunities for structured networking with peers working in similar domains and access to a resource library of code samples, case studies, and best practice guidelines.', + tags: ['tech', 'virtual', 'networking', 'learning', 'AI', 'ethics', 'development', 'workshops'] }, { id: 'event2', type: 'event', noun: 'event', - name: 'International Food Festival', - content: 'Annual celebration of global cuisines featuring cooking demonstrations, tastings, and cultural performances.', - tags: ['food', 'festival', 'international', 'culture'] + name: 'International Food Festival: Culinary Traditions in a Modern World', + content: 'Annual celebration of global cuisines featuring cooking demonstrations from renowned chefs representing diverse culinary traditions. The festival includes interactive tastings with detailed exploration of ingredient origins and cultural significance, panel discussions on preserving food heritage while embracing innovation, and cultural performances that contextualize food within broader cultural expressions. Special exhibition areas focus on sustainable food systems, the impact of climate change on traditional agriculture, and the role of technology in modern food production and distribution.', + tags: ['food', 'festival', 'international', 'culture', 'cooking', 'sustainability', 'traditions', 'innovation'] } ] - // Add Nouns to the database - for (const item of sampleNouns) { - // Pass the content as the data to vectorize, and the entire item as metadata - await database.add(item.content || item.name, item) - } + // Add Nouns to the database using batch processing for better performance + const nounItems = sampleNouns.map(item => ({ + vectorOrData: item.content || item.name, + metadata: item + })) - // Wait a moment to ensure nouns are fully indexed before adding verbs - await new Promise(resolve => setTimeout(resolve, 500)) + await database.addBatch(nounItems, { batchSize: 50 }) + + // No need to wait as addBatch ensures all items are fully processed // Verify nouns exist before adding verbs const entities = await database.getAllNouns() @@ -1921,24 +1925,38 @@ await db.init() } ] - // Add Verbs to the database - for (const item of sampleVerbs) { - try { - // Check if source and target nouns exist before adding the verb - if (!nounIds.includes(item.source)) { - throw new Error(`Source noun with ID ${item.source} not found. Available nouns: ${nounIds.join(', ')}`) - } - if (!nounIds.includes(item.target)) { - throw new Error(`Target noun with ID ${item.target} not found. Available nouns: ${nounIds.join(', ')}`) - } - - await database.relate(item.source, item.target, item.verb, item.data) - console.log(`Added verb: ${item.source} ${item.verb} ${item.target}`) - } catch (verbError) { - console.error(`Failed to add verb: ${verbError.message}`) - log('db-status', `Failed to add verb: ${verbError.message}`) + // Add Verbs to the database using parallel processing for better performance + // First filter out verbs with invalid source or target nouns + const validVerbs = sampleVerbs.filter(item => { + if (!nounIds.includes(item.source)) { + console.error(`Source noun with ID ${item.source} not found. Available nouns: ${nounIds.join(', ')}`) + log('db-status', `Failed to add verb: Source noun with ID ${item.source} not found`) + return false } - } + if (!nounIds.includes(item.target)) { + console.error(`Target noun with ID ${item.target} not found. Available nouns: ${nounIds.join(', ')}`) + log('db-status', `Failed to add verb: Target noun with ID ${item.target} not found`) + return false + } + return true + }) + + // Process verbs in parallel using Promise.all + const verbPromises = validVerbs.map(item => + database.relate(item.source, item.target, item.verb, item.data) + .then(() => { + console.log(`Added verb: ${item.source} ${item.verb} ${item.target}`) + return true + }) + .catch(verbError => { + console.error(`Failed to add verb: ${verbError.message}`) + log('db-status', `Failed to add verb: ${verbError.message}`) + return false + }) + ) + + // Wait for all verbs to be processed + await Promise.all(verbPromises) // Count successfully added verbs const addedVerbs = await database.getAllVerbs() diff --git a/package.json b/package.json index 401e3334..5cdc4931 100644 --- a/package.json +++ b/package.json @@ -100,6 +100,7 @@ "dependencies": { "@aws-sdk/client-s3": "^3.427.0", "@tensorflow-models/universal-sentence-encoder": "^1.3.3", + "@tensorflow/tfjs": "^4.22.0", "@tensorflow/tfjs-backend-cpu": "^4.22.0", "@tensorflow/tfjs-backend-webgl": "^4.22.0", "@tensorflow/tfjs-converter": "^4.22.0", diff --git a/src/brainyData.ts b/src/brainyData.ts index b99606ee..b7b1f743 100644 --- a/src/brainyData.ts +++ b/src/brainyData.ts @@ -24,7 +24,9 @@ import { import { cosineDistance, defaultEmbeddingFunction, - euclideanDistance + defaultBatchEmbeddingFunction, + euclideanDistance, + cleanupWorkerPools } from './utils/index.js' import { NounType, VerbType, GraphNoun } from './types/graphTypes.js' import { @@ -127,6 +129,7 @@ export class BrainyData implements BrainyDataInterface { private index: HNSWIndex | HNSWIndexOptimized private storage: StorageAdapter | null = null private isInitialized = false + private isInitializing = false private embeddingFunction: EmbeddingFunction private distanceFunction: DistanceFunction private requestPersistentStorage: boolean @@ -204,7 +207,48 @@ export class BrainyData implements BrainyDataInterface { return } + // Prevent recursive initialization + if (this.isInitializing) { + return + } + + this.isInitializing = true + try { + // Pre-load the embedding model early to ensure it's always available + // This helps prevent issues with the Universal Sentence Encoder not being loaded + try { + console.log('Pre-loading Universal Sentence Encoder model...') + // Call embedding function directly to avoid circular dependency with embed() + await this.embeddingFunction('') + console.log('Universal Sentence Encoder model loaded successfully') + } catch (embedError) { + console.warn('Failed to pre-load Universal Sentence Encoder:', embedError) + + // Try again with a retry mechanism + console.log('Retrying Universal Sentence Encoder initialization...') + try { + // Wait a moment before retrying + await new Promise(resolve => setTimeout(resolve, 1000)) + + // Try again with a different approach - use the non-threaded version + // This is a fallback in case the threaded version fails + const { createTensorFlowEmbeddingFunction } = await import('./utils/embedding.js') + const fallbackEmbeddingFunction = createTensorFlowEmbeddingFunction() + + // Test the fallback embedding function + await fallbackEmbeddingFunction('') + + // If successful, replace the embedding function + console.log('Successfully loaded Universal Sentence Encoder with fallback method') + this.embeddingFunction = fallbackEmbeddingFunction + } catch (retryError) { + console.error('All attempts to load Universal Sentence Encoder failed:', retryError) + // Continue initialization even if embedding model fails to load + // The application will need to handle missing embedding functionality + } + } + // Initialize storage if not provided in constructor if (!this.storage) { // Combine storage config with requestPersistentStorage for backward compatibility @@ -251,8 +295,10 @@ export class BrainyData implements BrainyDataInterface { } this.isInitialized = true + this.isInitializing = false } catch (error) { console.error('Failed to initialize BrainyData:', error) + this.isInitializing = false throw new Error(`Failed to initialize BrainyData: ${error}`) } } @@ -478,6 +524,7 @@ export class BrainyData implements BrainyDataInterface { forceEmbed?: boolean // Force using the embedding function even if input is a vector addToRemote?: boolean // Whether to also add to the remote server if connected concurrency?: number // Maximum number of concurrent operations (default: 4) + batchSize?: number // Maximum number of items to process in a single batch (default: 50) } = {} ): Promise { await this.ensureInitialized() @@ -488,22 +535,86 @@ export class BrainyData implements BrainyDataInterface { // Default concurrency to 4 if not specified const concurrency = options.concurrency || 4 + // Default batch size to 50 if not specified + const batchSize = options.batchSize || 50 + try { - // Process items in batches to control concurrency + // Process items in batches to control concurrency and memory usage const ids: string[] = [] const itemsToProcess = [...items] // Create a copy to avoid modifying the original array while (itemsToProcess.length > 0) { - // Take up to 'concurrency' items to process in parallel - const batch = itemsToProcess.splice(0, concurrency) + // Take up to 'batchSize' items to process in a batch + const batch = itemsToProcess.splice(0, batchSize) - // Process this batch in parallel - const batchResults = await Promise.all( - batch.map((item) => - this.add(item.vectorOrData, item.metadata, options) - ) + // Separate items that are already vectors from those that need embedding + const vectorItems: Array<{ + vectorOrData: Vector + metadata?: T + index: number + }> = [] + + const textItems: Array<{ + text: string + metadata?: T + index: number + }> = [] + + // Categorize items + batch.forEach((item, index) => { + if ( + Array.isArray(item.vectorOrData) && + item.vectorOrData.every((val) => typeof val === 'number') && + !options.forceEmbed + ) { + // Item is already a vector + vectorItems.push({ + vectorOrData: item.vectorOrData, + metadata: item.metadata, + index + }) + } else if (typeof item.vectorOrData === 'string') { + // Item is text that needs embedding + textItems.push({ + text: item.vectorOrData, + metadata: item.metadata, + index + }) + } else { + // For now, treat other types as text + // In a more complete implementation, we might handle other types differently + const textRepresentation = String(item.vectorOrData) + textItems.push({ + text: textRepresentation, + metadata: item.metadata, + index + }) + } + }) + + // Process vector items (already embedded) + const vectorPromises = vectorItems.map(item => + this.add(item.vectorOrData, item.metadata, options) ) + // Process text items in a single batch embedding operation + let textPromises: Promise[] = [] + if (textItems.length > 0) { + // Extract just the text for batch embedding + const texts = textItems.map(item => item.text) + + // Perform batch embedding + const embeddings = await defaultBatchEmbeddingFunction(texts) + + // Add each item with its embedding + textPromises = textItems.map((item, i) => + this.add(embeddings[i], item.metadata, { ...options, forceEmbed: false }) + ) + } + + // Combine all promises + const batchResults = await Promise.all([...vectorPromises, ...textPromises]) + // Add the results to our ids array ids.push(...batchResults) } @@ -1729,7 +1840,28 @@ export class BrainyData implements BrainyDataInterface { * Ensure the database is initialized */ private async ensureInitialized(): Promise { - if (!this.isInitialized) { + if (this.isInitialized) { + return + } + + if (this.isInitializing) { + // If initialization is already in progress, wait for it to complete + // by polling the isInitialized flag + let attempts = 0 + const maxAttempts = 100 // Prevent infinite loop + const delay = 50 // ms + + while (this.isInitializing && !this.isInitialized && attempts < maxAttempts) { + await new Promise(resolve => setTimeout(resolve, delay)) + attempts++ + } + + if (!this.isInitialized) { + // If still not initialized after waiting, try to initialize again + await this.init() + } + } else { + // Normal case - not initialized and not initializing await this.init() } } @@ -1838,6 +1970,9 @@ export class BrainyData implements BrainyDataInterface { await this.disconnectFromRemoteServer() } + // Clean up worker pools to release resources + cleanupWorkerPools() + // Additional cleanup could be added here in the future this.isInitialized = false diff --git a/src/global.d.ts b/src/global.d.ts index 5b172476..2256dbeb 100644 --- a/src/global.d.ts +++ b/src/global.d.ts @@ -5,11 +5,16 @@ // Extend the globalThis interface to include the __ENV__ property declare global { var __ENV__: { - isBrowser: boolean; - isNode: string | false; - isServerless: boolean; - }; + isBrowser: boolean + isNode: string | false + isServerless: boolean + } + + // Window interface extensions (worker-specific functions removed) + interface Window { + importTensorFlow?: () => Promise + } } // This export is needed to make this file a module -export {}; +export {} diff --git a/src/hnsw/hnswIndex.ts b/src/hnsw/hnswIndex.ts index c772e0ee..2585dac9 100644 --- a/src/hnsw/hnswIndex.ts +++ b/src/hnsw/hnswIndex.ts @@ -10,7 +10,7 @@ import { Vector, VectorDocument } from '../coreTypes.js' -import { euclideanDistance, calculateDistancesWithGPU } from '../utils/index.js' +import { euclideanDistance, calculateDistancesBatch } from '../utils/index.js' import { executeInThread } from '../utils/workerUtils.js' // Default HNSW parameters @@ -83,7 +83,7 @@ export class HNSWIndex { const vectorsOnly = vectors.map((item) => item.vector) // Use GPU-accelerated distance calculation when possible - const distances = await calculateDistancesWithGPU( + const distances = await calculateDistancesBatch( queryVector, vectorsOnly, this.distanceFunction @@ -96,51 +96,15 @@ export class HNSWIndex { })) } catch (error) { console.error( - 'Error in GPU-accelerated distance calculation, falling back to threaded CPU:', + 'Error in GPU-accelerated distance calculation, falling back to sequential processing:', error ) - // Fall back to threaded CPU processing if GPU acceleration fails - // Function to be executed in a worker thread - const distanceCalculator = (args: { - queryVector: Vector - vectors: Array<{ id: string; vector: Vector }> - distanceFnString: string - }) => { - const { queryVector, vectors, distanceFnString } = args - - // Recreate the distance function from its string representation - const distanceFunction = new Function( - 'return ' + distanceFnString - )() as DistanceFunction - - // Calculate distances for all items - return vectors.map((item) => ({ - id: item.id, - distance: distanceFunction(queryVector, item.vector) - })) - } - - try { - // Convert the distance function to a string for serialization - const distanceFnString = this.distanceFunction.toString() - - // Execute the distance calculation in a separate thread - return await executeInThread>( - distanceCalculator.toString(), - { queryVector, vectors, distanceFnString } - ) - } catch (threadError) { - console.error( - 'Error in threaded distance calculation, falling back to sequential:', - threadError - ) - // Fall back to sequential processing if both GPU and threaded execution fail - return vectors.map((item) => ({ - id: item.id, - distance: this.distanceFunction(queryVector, item.vector) - })) - } + // Fall back to sequential processing if GPU acceleration fails + return vectors.map((item) => ({ + id: item.id, + distance: this.distanceFunction(queryVector, item.vector) + })) } } diff --git a/src/utils/distance.ts b/src/utils/distance.ts index 8931d7fd..90917dd2 100644 --- a/src/utils/distance.ts +++ b/src/utils/distance.ts @@ -13,7 +13,10 @@ import { isThreadingAvailable } from './environment.js' * Lower values indicate higher similarity * Optimized using array methods for Node.js 23.11+ */ -export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): number => { +export const euclideanDistance: DistanceFunction = ( + a: Vector, + b: Vector +): number => { if (a.length !== b.length) { throw new Error('Vectors must have the same dimensions') } @@ -21,7 +24,7 @@ export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): numbe // Use array.reduce for better performance in Node.js 23.11+ const sum = a.reduce((acc, val, i) => { const diff = val - b[i] - return acc + (diff * diff) + return acc + diff * diff }, 0) return Math.sqrt(sum) @@ -33,19 +36,25 @@ export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): numbe * Range: 0 (identical) to 2 (opposite) * Optimized using array methods for Node.js 23.11+ */ -export const cosineDistance: DistanceFunction = (a: Vector, b: Vector): number => { +export const cosineDistance: DistanceFunction = ( + a: Vector, + b: Vector +): number => { if (a.length !== b.length) { throw new Error('Vectors must have the same dimensions') } // Use array.reduce to calculate all values in a single pass - const { dotProduct, normA, normB } = a.reduce((acc, val, i) => { - return { - dotProduct: acc.dotProduct + (val * b[i]), - normA: acc.normA + (val * val), - normB: acc.normB + (b[i] * b[i]) - } - }, { dotProduct: 0, normA: 0, normB: 0 }) + const { dotProduct, normA, normB } = a.reduce( + (acc, val, i) => { + return { + dotProduct: acc.dotProduct + val * b[i], + normA: acc.normA + val * val, + normB: acc.normB + b[i] * b[i] + } + }, + { dotProduct: 0, normA: 0, normB: 0 } + ) if (normA === 0 || normB === 0) { return 2 // Maximum distance for zero vectors @@ -61,7 +70,10 @@ export const cosineDistance: DistanceFunction = (a: Vector, b: Vector): number = * Lower values indicate higher similarity * Optimized using array methods for Node.js 23.11+ */ -export const manhattanDistance: DistanceFunction = (a: Vector, b: Vector): number => { +export const manhattanDistance: DistanceFunction = ( + a: Vector, + b: Vector +): number => { if (a.length !== b.length) { throw new Error('Vectors must have the same dimensions') } @@ -76,209 +88,211 @@ export const manhattanDistance: DistanceFunction = (a: Vector, b: Vector): numbe * Converted to a distance metric (lower is better) * Optimized using array methods for Node.js 23.11+ */ -export const dotProductDistance: DistanceFunction = (a: Vector, b: Vector): number => { +export const dotProductDistance: DistanceFunction = ( + a: Vector, + b: Vector +): number => { if (a.length !== b.length) { throw new Error('Vectors must have the same dimensions') } // Use array.reduce for better performance in Node.js 23.11+ - const dotProduct = a.reduce((sum, val, i) => sum + (val * b[i]), 0) + const dotProduct = a.reduce((sum, val, i) => sum + val * b[i], 0) // Convert to a distance metric (lower is better) return -dotProduct } /** - * GPU-accelerated batch distance calculation - * Uses TensorFlow.js with WebGL backend when available for optimal performance - * Falls back to CPU processing when GPU is not available + * Batch distance calculation + * Uses TensorFlow.js with CPU backend for optimized performance * * @param queryVector The query vector to compare against all vectors * @param vectors Array of vectors to compare against * @param distanceFunction The distance function to use * @returns Promise resolving to array of distances */ -export async function calculateDistancesWithGPU( +export async function calculateDistancesBatch( queryVector: Vector, vectors: Vector[], distanceFunction: DistanceFunction = euclideanDistance ): Promise { // For small batches, use the standard distance function if (vectors.length < 10) { - return vectors.map(vector => distanceFunction(queryVector, vector)) + return vectors.map((vector) => distanceFunction(queryVector, vector)) } try { // Function to be executed in a worker thread - const distanceCalculator = async ( - args: { - queryVector: Vector, - vectors: Vector[], - distanceFnString: string - } - ) => { + const distanceCalculator = async (args: { + queryVector: Vector + vectors: Vector[] + distanceFnString: string + }) => { const { queryVector, vectors, distanceFnString } = args - // Try to use TensorFlow.js with GPU acceleration if available + // Use TensorFlow.js with CPU processing const useTensorFlow = async () => { - try { - // TensorFlow.js will use its default EPSILON value + // TensorFlow.js will use its default EPSILON value + // Use the importTensorFlow function if available (in worker context) + // or directly import TensorFlow.js (in main thread) + let tf + if ( + typeof self !== 'undefined' && + typeof self.importTensorFlow === 'function' + ) { + // In worker context, use the importTensorFlow function + tf = await self.importTensorFlow() + } else { // Dynamically import TensorFlow.js core module and backends - const tf = await import('@tensorflow/tfjs-core') + tf = await import('@tensorflow/tfjs-core') - // Import CPU backend as fallback + // Import CPU backend await import('@tensorflow/tfjs-backend-cpu') - let usingGPU = false + // Set CPU as the backend + await tf.setBackend('cpu') + } - try { - // Try to import and use WebGL backend (GPU) - await import('@tensorflow/tfjs-backend-webgl') + // Convert vectors to tensors + const queryTensor = tf.tensor2d([queryVector]) + const vectorsTensor = tf.tensor2d(vectors) - // Check if WebGL is available and set it as the backend - if (await tf.findBackend('webgl') || await tf.ready().then(() => tf.findBackend('webgl'))) { - await tf.setBackend('webgl') - usingGPU = true - } else { - await tf.setBackend('cpu') - } - } catch (err) { - // If WebGL fails, use CPU - await tf.setBackend('cpu') - } + let distances: number[] - // Convert vectors to tensors - const queryTensor = tf.tensor2d([queryVector]) - const vectorsTensor = tf.tensor2d(vectors) + // Calculate distances based on the distance function type + if (distanceFnString.includes('euclideanDistance')) { + // Euclidean distance using GPU-optimized operations + // Formula: sqrt(sum((a - b)^2)) + const expanded = tf.sub( + (queryTensor as any).expandDims(1), + (vectorsTensor as any).expandDims(0) + ) + const squaredDiff = tf.square(expanded) + const sumSquaredDiff = tf.sum(squaredDiff, -1) + const distancesTensor = tf.sqrt(sumSquaredDiff) + distances = (await (distancesTensor as any) + .squeeze() + .array()) as number[] - let distances: number[] + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + expanded.dispose() + squaredDiff.dispose() + sumSquaredDiff.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('cosineDistance')) { + // Cosine distance using GPU-optimized operations + // Formula: 1 - (a·b / (||a|| * ||b||)) + const dotProduct = tf.matMul( + queryTensor, + (vectorsTensor as any).transpose() + ) - // Calculate distances based on the distance function type - if (distanceFnString.includes('euclideanDistance')) { - // Euclidean distance using GPU-optimized operations - // Formula: sqrt(sum((a - b)^2)) - const expanded = tf.sub((queryTensor as any).expandDims(1), (vectorsTensor as any).expandDims(0)) - const squaredDiff = tf.square(expanded) - const sumSquaredDiff = tf.sum(squaredDiff, -1) - const distancesTensor = tf.sqrt(sumSquaredDiff) - distances = await (distancesTensor as any).squeeze().array() as number[] + const queryNorm = tf.norm(queryTensor, 2, 1) + const vectorsNorm = tf.norm(vectorsTensor, 2, 1) - // Clean up tensors - queryTensor.dispose() - vectorsTensor.dispose() - expanded.dispose() - squaredDiff.dispose() - sumSquaredDiff.dispose() - distancesTensor.dispose() - } else if (distanceFnString.includes('cosineDistance')) { - // Cosine distance using GPU-optimized operations - // Formula: 1 - (a·b / (||a|| * ||b||)) - const dotProduct = tf.matMul(queryTensor, (vectorsTensor as any).transpose()) + const normProduct = tf.outerProduct( + queryNorm as any, + vectorsNorm as any + ) + const cosineSimilarity = tf.div(dotProduct, normProduct) + const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity) - const queryNorm = tf.norm(queryTensor, 2, 1) - const vectorsNorm = tf.norm(vectorsTensor, 2, 1) + distances = (await (distancesTensor as any) + .squeeze() + .array()) as number[] - const normProduct = tf.outerProduct(queryNorm as any, vectorsNorm as any) - const cosineSimilarity = tf.div(dotProduct, normProduct) - const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity) + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + dotProduct.dispose() + queryNorm.dispose() + vectorsNorm.dispose() + normProduct.dispose() + cosineSimilarity.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('manhattanDistance')) { + // Manhattan distance using GPU-optimized operations + // Formula: sum(|a - b|) + const diff = tf.sub( + (queryTensor as any).expandDims(1), + (vectorsTensor as any).expandDims(0) + ) + const absDiff = tf.abs(diff) + const distancesTensor = tf.sum(absDiff, -1) - distances = await (distancesTensor as any).squeeze().array() as number[] + distances = (await (distancesTensor as any) + .squeeze() + .array()) as number[] - // Clean up tensors - queryTensor.dispose() - vectorsTensor.dispose() - dotProduct.dispose() - queryNorm.dispose() - vectorsNorm.dispose() - normProduct.dispose() - cosineSimilarity.dispose() - distancesTensor.dispose() - } else if (distanceFnString.includes('manhattanDistance')) { - // Manhattan distance using GPU-optimized operations - // Formula: sum(|a - b|) - const diff = tf.sub((queryTensor as any).expandDims(1), (vectorsTensor as any).expandDims(0)) - const absDiff = tf.abs(diff) - const distancesTensor = tf.sum(absDiff, -1) + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + diff.dispose() + absDiff.dispose() + distancesTensor.dispose() + } else if (distanceFnString.includes('dotProductDistance')) { + // Dot product distance using GPU-optimized operations + // Formula: -sum(a * b) + const dotProduct = tf.matMul( + queryTensor, + (vectorsTensor as any).transpose() + ) + const distancesTensor = tf.neg(dotProduct) - distances = await (distancesTensor as any).squeeze().array() as number[] + distances = (await (distancesTensor as any) + .squeeze() + .array()) as number[] - // Clean up tensors - queryTensor.dispose() - vectorsTensor.dispose() - diff.dispose() - absDiff.dispose() - distancesTensor.dispose() - } else if (distanceFnString.includes('dotProductDistance')) { - // Dot product distance using GPU-optimized operations - // Formula: -sum(a * b) - const dotProduct = tf.matMul(queryTensor, (vectorsTensor as any).transpose()) - const distancesTensor = tf.neg(dotProduct) + // Clean up tensors + queryTensor.dispose() + vectorsTensor.dispose() + dotProduct.dispose() + distancesTensor.dispose() + } else { + // For unknown distance functions, fall back to direct CPU implementation + throw new Error( + 'Unsupported distance function for TensorFlow optimization' + ) + } - distances = await (distancesTensor as any).squeeze().array() as number[] - - // Clean up tensors - queryTensor.dispose() - vectorsTensor.dispose() - dotProduct.dispose() - distancesTensor.dispose() - } else { - // For unknown distance functions, fall back to CPU implementation - throw new Error('Unsupported distance function for GPU acceleration') - } - - return { - distances, - usingGPU - } - } catch (error) { - // If TensorFlow.js fails, fall back to CPU implementation - throw error + return { + distances } } - // Try to use TensorFlow.js with GPU acceleration + // Try to use TensorFlow.js with CPU optimization try { return await useTensorFlow() } catch (error) { - // Fall back to CPU implementation if TensorFlow.js fails + // Fall back to direct CPU implementation if TensorFlow.js fails // Recreate the distance function from its string representation - const distanceFunction = new Function('return ' + distanceFnString)() as DistanceFunction + const distanceFunction = new Function( + 'return ' + distanceFnString + )() as DistanceFunction // Calculate distances for all vectors - const distances = vectors.map(vector => distanceFunction(queryVector, vector)) + const distances = vectors.map((vector) => + distanceFunction(queryVector, vector) + ) return { - distances, - usingGPU: false + distances } } } - // Execute the distance calculation in a separate thread if threading is available - if (isThreadingAvailable()) { - try { - // Convert the distance function to a string for serialization - const distanceFnString = distanceFunction.toString() - - // Execute in a separate thread - const result = await executeInThread<{ distances: number[], usingGPU: boolean }>( - distanceCalculator.toString(), - { queryVector, vectors, distanceFnString } - ) - - return result.distances - } catch (error) { - // Fall back to main thread if threading fails - console.warn('Threaded distance calculation failed, falling back to main thread:', error) - } - } + // Threading is not available, so we'll always use the main thread implementation + // This comment is kept for clarity about the removed code // If threading is not available or failed, calculate distances in the main thread - return vectors.map(vector => distanceFunction(queryVector, vector)) + return vectors.map((vector) => distanceFunction(queryVector, vector)) } catch (error) { // If anything fails, fall back to the standard distance function - console.error('GPU-accelerated distance calculation failed:', error) - return vectors.map(vector => distanceFunction(queryVector, vector)) + console.error('Batch distance calculation failed:', error) + return vectors.map((vector) => distanceFunction(queryVector, vector)) } } diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index 4c8c8c46..1c130869 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -10,14 +10,15 @@ import { executeInThread } from './workerUtils.js' * This model provides high-quality text embeddings using TensorFlow.js * The required TensorFlow.js dependencies are automatically installed with this package * - * This implementation will use GPU acceleration via WebGL when available, - * falling back to CPU processing when GPU is not available or fails to initialize. + * This implementation attempts to use GPU processing when available for better performance, + * falling back to CPU processing for compatibility across all environments. */ export class UniversalSentenceEncoder implements EmbeddingModel { private model: any = null private initialized = false private tf: any = null private use: any = null + private backend: string = 'cpu' // Default to CPU /** * Initialize the embedding model @@ -47,36 +48,59 @@ export class UniversalSentenceEncoder implements EmbeddingModel { // Use type assertions to tell TypeScript these modules exist this.tf = await import('@tensorflow/tfjs-core') - // Import CPU and WebGL backends + // Import CPU backend (always needed as fallback) await import('@tensorflow/tfjs-backend-cpu') + // Try to import WebGL backend for GPU acceleration in browser environments try { - // Try to import and use WebGL backend first (GPU) - await import('@tensorflow/tfjs-backend-webgl') - - // Check if WebGL is available and set it as the backend - if ( - (await this.tf.findBackend('webgl')) || - (await this.tf.ready().then(() => this.tf.findBackend('webgl'))) - ) { - console.log('Using WebGL backend (GPU acceleration)') - await this.tf.setBackend('webgl') - } else { - console.log('WebGL backend not available, falling back to CPU') - await this.tf.setBackend('cpu') + if (typeof window !== 'undefined') { + await import('@tensorflow/tfjs-backend-webgl') + // Check if WebGL is available using setBackend instead of findBackend + try { + if (this.tf.setBackend) { + await this.tf.setBackend('webgl') + this.backend = 'webgl' + console.log('Using WebGL backend for TensorFlow.js') + } else { + console.warn( + 'tf.setBackend is not available, falling back to CPU' + ) + } + } catch (e) { + console.warn('WebGL backend not available, falling back to CPU:', e) + this.backend = 'cpu' + } } - } catch (err) { - console.warn( - 'WebGL backend failed to initialize, using CPU backend:', - err - ) - await this.tf.setBackend('cpu') + } catch (error) { + console.warn('WebGL backend not available, falling back to CPU:', error) + this.backend = 'cpu' + } + + // Set the backend + if (this.tf.setBackend) { + await this.tf.setBackend(this.backend) } this.use = await import('@tensorflow-models/universal-sentence-encoder') + // Log the module structure to help with debugging + console.log( + 'Universal Sentence Encoder module structure in main thread:', + Object.keys(this.use), + this.use.default ? Object.keys(this.use.default) : 'No default export' + ) + + // Try to find the load function in different possible module structures + const loadFunction = findUSELoadFunction(this.use) + + if (!loadFunction) { + throw new Error( + 'Could not find Universal Sentence Encoder load function' + ) + } + // Load the model - this.model = await this.use.load() + this.model = await loadFunction() this.initialized = true // Restore original console.warn @@ -132,6 +156,10 @@ export class UniversalSentenceEncoder implements EmbeddingModel { // Convert to array and return the first embedding const embeddingArray = await embeddings.array() + + // Dispose of the tensor to free memory + embeddings.dispose() + return embeddingArray[0] } catch (error) { console.error( @@ -144,6 +172,70 @@ export class UniversalSentenceEncoder implements EmbeddingModel { } } + /** + * Embed multiple texts into vectors using Universal Sentence Encoder + * This is more efficient than calling embed() multiple times + * @param dataArray Array of texts to embed + * @returns Array of embedding vectors + */ + public async embedBatch(dataArray: string[]): Promise { + if (!this.initialized) { + await this.init() + } + + try { + // Handle empty array case + if (dataArray.length === 0) { + return [] + } + + // Filter out empty strings and handle edge cases + const textToEmbed = dataArray.filter( + (text: string) => typeof text === 'string' && text.trim() !== '' + ) + + // If all strings were empty, return appropriate zero vectors + if (textToEmbed.length === 0) { + return dataArray.map(() => new Array(512).fill(0)) + } + + // Get embeddings for all texts in a single batch operation + const embeddings = await this.model.embed(textToEmbed) + + // Convert to array + const embeddingArray = await embeddings.array() + + // Dispose of the tensor to free memory + embeddings.dispose() + + // Map the results back to the original array order + const results: Vector[] = [] + let embeddingIndex = 0 + + for (let i = 0; i < dataArray.length; i++) { + const text = dataArray[i] + if (typeof text === 'string' && text.trim() !== '') { + // Use the embedding for non-empty strings + results.push(embeddingArray[embeddingIndex]) + embeddingIndex++ + } else { + // Use a zero vector for empty strings + results.push(new Array(512).fill(0)) + } + } + + return results + } catch (error) { + console.error( + 'Failed to batch embed text with Universal Sentence Encoder:', + error + ) + throw new Error( + `Failed to batch embed text with Universal Sentence Encoder: ${error}` + ) + } + } + /** * Dispose of the model resources */ @@ -162,6 +254,105 @@ export class UniversalSentenceEncoder implements EmbeddingModel { } } +/** + * Helper function to load the Universal Sentence Encoder model + * This tries multiple approaches to find the correct load function + * @param sentenceEncoderModule The imported module + * @returns The load function or null if not found + */ +function findUSELoadFunction(sentenceEncoderModule: any): Function | null { + // Log the module structure for debugging + console.log( + 'Universal Sentence Encoder module structure:', + Object.keys(sentenceEncoderModule), + sentenceEncoderModule.default + ? Object.keys(sentenceEncoderModule.default) + : 'No default export' + ) + + let loadFunction = null + + // Try sentenceEncoderModule.load first (direct export) + if ( + sentenceEncoderModule.load && + typeof sentenceEncoderModule.load === 'function' + ) { + loadFunction = sentenceEncoderModule.load + console.log('Using sentenceEncoderModule.load') + } + // Then try sentenceEncoderModule.default.load (default export) + else if ( + sentenceEncoderModule.default && + sentenceEncoderModule.default.load && + typeof sentenceEncoderModule.default.load === 'function' + ) { + loadFunction = sentenceEncoderModule.default.load + console.log('Using sentenceEncoderModule.default.load') + } + // Try sentenceEncoderModule.default directly if it's a function + else if ( + sentenceEncoderModule.default && + typeof sentenceEncoderModule.default === 'function' + ) { + loadFunction = sentenceEncoderModule.default + console.log('Using sentenceEncoderModule.default as function') + } + // Try sentenceEncoderModule directly if it's a function + else if (typeof sentenceEncoderModule === 'function') { + loadFunction = sentenceEncoderModule + console.log('Using sentenceEncoderModule as function') + } + // Try additional common patterns + else if ( + sentenceEncoderModule.UniversalSentenceEncoder && + typeof sentenceEncoderModule.UniversalSentenceEncoder.load === 'function' + ) { + loadFunction = sentenceEncoderModule.UniversalSentenceEncoder.load + console.log('Using sentenceEncoderModule.UniversalSentenceEncoder.load') + } else if ( + sentenceEncoderModule.default && + sentenceEncoderModule.default.UniversalSentenceEncoder && + typeof sentenceEncoderModule.default.UniversalSentenceEncoder.load === + 'function' + ) { + loadFunction = sentenceEncoderModule.default.UniversalSentenceEncoder.load + console.log( + 'Using sentenceEncoderModule.default.UniversalSentenceEncoder.load' + ) + } + // Try to find the load function in the module's properties + else { + // Look for any property that might be a load function + for (const key in sentenceEncoderModule) { + if (typeof sentenceEncoderModule[key] === 'function') { + // Check if the function name or key contains 'load' + const fnName = sentenceEncoderModule[key].name || key; + if (fnName.toLowerCase().includes('load')) { + loadFunction = sentenceEncoderModule[key]; + console.log(`Using sentenceEncoderModule.${key} as load function`); + break; + } + } + // Also check nested objects + else if (typeof sentenceEncoderModule[key] === 'object' && sentenceEncoderModule[key] !== null) { + for (const nestedKey in sentenceEncoderModule[key]) { + if (typeof sentenceEncoderModule[key][nestedKey] === 'function') { + const fnName = sentenceEncoderModule[key][nestedKey].name || nestedKey; + if (fnName.toLowerCase().includes('load')) { + loadFunction = sentenceEncoderModule[key][nestedKey]; + console.log(`Using sentenceEncoderModule.${key}.${nestedKey} as load function`); + break; + } + } + } + if (loadFunction) break; + } + } + } + + return loadFunction; +} + /** * Create an embedding function from an embedding model * @param model Embedding model to use @@ -177,21 +368,22 @@ export function createEmbeddingFunction( /** * Creates a TensorFlow-based Universal Sentence Encoder embedding function * This is the required embedding function for all text embeddings + * Uses a shared model instance for better performance across multiple calls */ -export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { - // Create a single shared instance of the model - const model = new UniversalSentenceEncoder() - let modelInitialized = false +// Create a single shared instance of the model that persists across all embedding calls +const sharedModel = new UniversalSentenceEncoder() +let sharedModelInitialized = false +export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { return async (data: any): Promise => { try { // Initialize the model if it hasn't been initialized yet - if (!modelInitialized) { - await model.init() - modelInitialized = true + if (!sharedModelInitialized) { + await sharedModel.init() + sharedModelInitialized = true } - return await model.embed(data) + return await sharedModel.embed(data) } catch (error) { console.error('Failed to use TensorFlow embedding:', error) throw new Error( @@ -202,165 +394,59 @@ export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { } /** - * Creates a TensorFlow-based Universal Sentence Encoder embedding function that runs in a separate thread - * This provides better performance for embedding operations by: - * 1. Using GPU acceleration via WebGL when available - * 2. Running in a separate thread to avoid blocking the main thread - * 3. Falling back to CPU processing when GPU is not available - * - * @param options Configuration options - * @returns An embedding function that runs in a separate thread with GPU acceleration when available + * Default embedding function + * Uses UniversalSentenceEncoder for all text embeddings + * TensorFlow.js is required for this to work + * Uses CPU for compatibility */ -export function createThreadedEmbeddingFunction( - options: { fallbackToMain?: boolean } = {} -): EmbeddingFunction { - // Create a standard embedding function to use as fallback - const standardEmbedding = createTensorFlowEmbeddingFunction() +export const defaultEmbeddingFunction: EmbeddingFunction = + createTensorFlowEmbeddingFunction() - // Flag to track if we've fallen back to main thread - let useFallback = false +/** + * Default batch embedding function + * Uses UniversalSentenceEncoder for all text embeddings + * TensorFlow.js is required for this to work + * Processes all items in a single batch operation + * Uses a shared model instance for better performance across multiple calls + */ +// Create a single shared instance of the model that persists across function calls +const sharedBatchModel = new UniversalSentenceEncoder() +let sharedBatchModelInitialized = false - return async (data: any): Promise => { - // If we've already determined that threading doesn't work, use the fallback - if (useFallback) { - return standardEmbedding(data) +export const defaultBatchEmbeddingFunction: ( + dataArray: string[] +) => Promise = async (dataArray: string[]): Promise => { + try { + // Initialize the model if it hasn't been initialized yet + if (!sharedBatchModelInitialized) { + await sharedBatchModel.init() + sharedBatchModelInitialized = true } - try { - // Function to be executed in a worker thread - // This must be a regular function (not async) to avoid Promise cloning issues - const embedInWorker = (inputData: any) => { - // Return a plain object with the input data - // All async operations will be performed inside the worker - return { data: inputData } - } - - // Worker implementation function that will be stringified and run in the worker - const workerImplementation = async ({ data }: { data: any }) => { - try { - // We need to dynamically import TensorFlow.js core module and USE in the worker - // Use a variable name that won't conflict with any minified variables - - // TensorFlow.js will use its default EPSILON value - - // Import TensorFlow.js modules - const tf = await import('@tensorflow/tfjs-core') - - // Import CPU backend in the worker - await import('@tensorflow/tfjs-backend-cpu') - - try { - // Try to import and use WebGL backend first (GPU) - await import('@tensorflow/tfjs-backend-webgl') - - // Check if WebGL is available and set it as the backend - if ( - (tf.findBackend && (await tf.findBackend('webgl'))) || - (tf.ready && - (await tf - .ready() - .then(() => tf.findBackend && tf.findBackend('webgl')))) - ) { - console.log('Worker: Using WebGL backend (GPU acceleration)') - if (tf.setBackend) { - await tf.setBackend('webgl') - } - } else { - console.log( - 'Worker: WebGL backend not available, falling back to CPU' - ) - if (tf.setBackend) { - await tf.setBackend('cpu') - } - } - } catch (err) { - console.warn( - 'Worker: WebGL backend failed to initialize, using CPU backend:', - err - ) - if (tf.setBackend) { - await tf.setBackend('cpu') - } - } - - // Import the Universal Sentence Encoder - const sentenceEncoderModule = await import( - '@tensorflow-models/universal-sentence-encoder' - ) - - // Load the model directly from the module - const model = await sentenceEncoderModule.load() - - // Handle different input types - let textToEmbed: string[] - if (typeof data === 'string') { - if (data.trim() === '') { - return new Array(512).fill(0) - } - textToEmbed = [data] - } else if ( - Array.isArray(data) && - data.every((item) => typeof item === 'string') - ) { - if (data.length === 0 || data.every((item) => item.trim() === '')) { - return new Array(512).fill(0) - } - textToEmbed = data.filter((item) => item.trim() !== '') - if (textToEmbed.length === 0) { - return new Array(512).fill(0) - } - } else { - throw new Error( - 'UniversalSentenceEncoder only supports string or string[] data' - ) - } - - // Get embeddings - const embeddings = await model.embed(textToEmbed) - - // Convert to array and return the first embedding - const embeddingArray = await embeddings.array() - - // Dispose of the tensor to free memory - embeddings.dispose() - - return embeddingArray[0] - } catch (error) { - console.error('Worker error:', error) - throw error - } - } - - // Execute the embedding function in a separate thread - // Pass the worker implementation as a string to avoid Promise cloning issues - return await executeInThread( - workerImplementation.toString(), - embedInWorker(data) - ) - } catch (error) { - // If threading fails and fallback is enabled, use the standard embedding function - if (options.fallbackToMain) { - console.warn( - 'Threaded embedding failed, falling back to main thread:', - error - ) - useFallback = true - return standardEmbedding(data) - } - - // Otherwise, propagate the error - throw new Error(`Threaded embedding failed: ${error}`) - } + return await sharedBatchModel.embedBatch(dataArray) + } catch (error) { + console.error('Failed to use TensorFlow batch embedding:', error) + throw new Error( + `Universal Sentence Encoder batch embedding failed: ${error}` + ) } } /** - * Default embedding function - * Uses UniversalSentenceEncoder for all text embeddings - * TensorFlow.js is required for this to work - * Uses GPU acceleration via WebGL when available for optimal performance - * Uses threading when available to avoid blocking the main thread - * Falls back to CPU processing when GPU is not available + * Creates an embedding function that runs in a separate thread + * This is a wrapper around createEmbeddingFunction that uses executeInThread + * @param model Embedding model to use */ -export const defaultEmbeddingFunction: EmbeddingFunction = - createThreadedEmbeddingFunction({ fallbackToMain: true }) +export function createThreadedEmbeddingFunction( + model: EmbeddingModel +): EmbeddingFunction { + const embeddingFunction = createEmbeddingFunction(model) + + return async (data: any): Promise => { + // Convert the embedding function to a string + const fnString = embeddingFunction.toString() + + // Execute the embedding function in a "thread" (main thread in this implementation) + return await executeInThread(fnString, data) + } +} diff --git a/src/utils/environment.ts b/src/utils/environment.ts index 477b07ca..82c701d7 100644 --- a/src/utils/environment.ts +++ b/src/utils/environment.ts @@ -39,7 +39,7 @@ export function areWebWorkersAvailable(): boolean { */ export function areWorkerThreadsAvailable(): boolean { if (!isNode()) return false; - + try { // Dynamic import to avoid errors in browser environments require('worker_threads'); @@ -51,7 +51,8 @@ export function areWorkerThreadsAvailable(): boolean { /** * Determine if threading is available in the current environment + * Always returns false since multithreading has been removed */ export function isThreadingAvailable(): boolean { - return areWebWorkersAvailable() || areWorkerThreadsAvailable(); + return false; } diff --git a/src/utils/index.ts b/src/utils/index.ts index 11db368e..3b8efae5 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -1,2 +1,3 @@ export * from './distance.js' export * from './embedding.js' +export * from './workerUtils.js' diff --git a/src/utils/workerUtils.ts b/src/utils/workerUtils.ts index abef3832..68b2e244 100644 --- a/src/utils/workerUtils.ts +++ b/src/utils/workerUtils.ts @@ -1,162 +1,30 @@ /** - * Utility functions for working with Web Workers and Worker Threads + * Utility functions for executing functions (without Web Workers or Worker Threads) + * This is a replacement for the original multithreaded implementation */ -import { isNode, isBrowser } from './environment.js' - /** - * Execute a function in a Web Worker (browser environment) - * - * @param fnString The function to execute as a string - * @param args The arguments to pass to the function - * @returns A promise that resolves with the result of the function - */ -export function executeInWebWorker(fnString: string, args: any): Promise { - return new Promise((resolve, reject) => { - // Create a blob URL for the worker script - const workerScript = ` - // Define global variables to prevent "is not defined" errors - // This ensures that any minified variable names used by bundlers are available - self.index = {} - self.index$1 = {} - self.index$2 = {} - self.universalSentenceEncoder_esm = {} - self.universalSentenceEncoder = {} - self.tfjs = {} - self.tfjs_core = {} - self.tfjs_backend_cpu = {} - self.tfjs_backend_webgl = {} - - // Additional variables that might be generated by bundlers - self.use = {} - self.tf = {} - self.sentenceEncoder = {} - - self.onmessage = async function(e) { - try { - const fn = ${fnString} - const result = await fn(e.data) - self.postMessage({ success: true, data: result }) - } catch (error) { - self.postMessage({ - success: false, - error: error instanceof Error ? error.message : String(error) - }) - } - } - ` - - const blob = new Blob([workerScript], { type: 'application/javascript' }) - const blobURL = URL.createObjectURL(blob) - - // Create a worker - const worker = new Worker(blobURL) - - // Set up message handling - worker.onmessage = function ( - e: MessageEvent<{ success: boolean; data: T; error?: string }> - ) { - URL.revokeObjectURL(blobURL) // Clean up - worker.terminate() // Terminate the worker - - if (e.data.success) { - resolve(e.data.data) - } else { - reject(new Error(e.data.error)) - } - } - - worker.onerror = function (error: ErrorEvent) { - URL.revokeObjectURL(blobURL) // Clean up - worker.terminate() // Terminate the worker - reject(error) - } - - // Start the worker - worker.postMessage(args) - }) -} - -/** - * Execute a function in a Worker Thread (Node.js environment) - * - * @param fnString The function to execute as a string - * @param args The arguments to pass to the function - * @returns A promise that resolves with the result of the function - */ -export function executeInWorkerThread( - fnString: string, - args: any -): Promise { - return new Promise((resolve, reject) => { - try { - // Dynamic import to avoid errors in browser environments - const { Worker } = require('worker_threads') - - // Create a worker script - const workerScript = ` - const { parentPort } = require('worker_threads') - - parentPort.once('message', async (data) => { - try { - const fn = ${fnString} - const result = await fn(data) - parentPort.postMessage({ success: true, data: result }) - } catch (error) { - parentPort.postMessage({ - success: false, - error: error instanceof Error ? error.message : String(error) - }) - } - }) - ` - - // Create a worker - const worker = new Worker(workerScript, { eval: true }) - - // Set up message handling - worker.on( - 'message', - (data: { success: boolean; data: T; error?: string }) => { - worker.terminate() // Terminate the worker - - if (data.success) { - resolve(data.data) - } else { - reject(new Error(data.error)) - } - } - ) - - worker.on('error', (error: Error) => { - worker.terminate() // Terminate the worker - reject(error) - }) - - // Start the worker - worker.postMessage(args) - } catch (error) { - reject(error) - } - }) -} - -/** - * Execute a function in a separate thread based on the environment + * Execute a function directly in the main thread * * @param fnString The function to execute as a string * @param args The arguments to pass to the function * @returns A promise that resolves with the result of the function */ export function executeInThread(fnString: string, args: any): Promise { - if (isBrowser()) { - return executeInWebWorker(fnString, args) - } else if (isNode()) { - return executeInWorkerThread(fnString, args) - } else { - // Fall back to executing in the main thread + try { // Parse the function from string and execute it const fn = new Function('return ' + fnString)() return Promise.resolve(fn(args) as T) + } catch (error) { + return Promise.reject(error) } } + +/** + * No-op function for backward compatibility + * This function does nothing since there are no worker pools to clean up + */ +export function cleanupWorkerPools(): void { + // No-op function + console.log('Worker pools cleanup called (no-op in non-threaded version)') +}