soulcraft/brainy
2
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions

feat(src/brainyData, src/utils): enhance embedding efficiency with batch processing and initialize safeguards

Browse source 
- 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.
This commit is contained in:
David Snelling 2025-06-27 14:06:59 -07:00
parent e81979dc84
commit bba9a0c219
11 changed files with 729 additions and 633 deletions
Download patch file Download diff file Expand all files Collapse all files

57
README.md
View file

@ -508,7 +508,7 @@ const id = await db.add(textOrVector, {
// other metadata... // other metadata...
}) })
// Add multiple nouns in parallel (with multithreading) // Add multiple nouns in parallel (with multithreading and batch embedding)
const ids = await db.addBatch([ const ids = await db.addBatch([
{ {
vectorOrData: "First item to add", vectorOrData: "First item to add",
@ -521,7 +521,8 @@ const ids = await db.addBatch([
// More items... // More items...
], { ], {
forceEmbed: false, 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 // Retrieve a noun
@ -591,9 +592,7 @@ await db.init()
// Or use the threaded embedding function for better performance // Or use the threaded embedding function for better performance
const threadedDb = new BrainyData({ const threadedDb = new BrainyData({
embeddingFunction: createThreadedEmbeddingFunction({ embeddingFunction: createThreadedEmbeddingFunction()
fallbackToMain: true // Fall back to main thread if threading fails
})
}) })
await threadedDb.init() 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 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 performance, especially for embedding operations. It uses GPU acceleration when available (via WebGL in browsers) and
threading is not available in the current environment. 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 ### 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 1. **GPU-Accelerated Embeddings**: Generate text embeddings using TensorFlow.js with WebGL backend when available
2. **GPU-Accelerated Distance Calculations**: Perform vector similarity calculations on the GPU for faster search 2. **Automatic Fallback**: Falls back to CPU backend when GPU is not available
3. **Automatic Fallback**: Falls back to CPU processing when GPU is not available or fails to initialize 3. **Optimized Distance Calculations**: Perform vector similarity calculations with optimized algorithms
4. **Cross-Environment Support**: Works in browsers (via WebGL) and Node.js environments 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 #### 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 1. **Parallel Batch Processing**: Add multiple items concurrently with controlled parallelism
2. **Multithreaded Vector Search**: Perform distance calculations in parallel for faster search operations 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 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 ```typescript
import { BrainyData, euclideanDistance } from '@soulcraft/brainy' import { BrainyData, euclideanDistance } from '@soulcraft/brainy'
@ -645,8 +650,6 @@ const db = new BrainyData({
// Performance optimization options // Performance optimization options
performance: { performance: {
useParallelization: true, // Enable multithreaded search operations 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 // 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 - `manhattanDistance`: Measures the sum of absolute differences between vector components
- `dotProductDistance`: Measures the negative dot product between vectors - `dotProductDistance`: Measures the negative dot product between vectors
All distance functions have GPU-accelerated versions that are automatically used when: All distance functions are optimized for performance and automatically use the most efficient implementation based on
1. GPU acceleration is enabled in the configuration the dataset size and available resources. For large datasets and high-dimensional vectors, Brainy uses batch processing
2. The operation involves a large number of vectors and multithreading when available to improve performance.
3. WebGL is available in the environment
The GPU-accelerated distance calculations provide significant performance improvements for large datasets and high-dimensional vectors.
## Backup and Restore ## Backup and Restore
@ -815,6 +815,9 @@ brainy import-sparse --input sparse-data.json
Brainy uses the following embedding approach: Brainy uses the following embedding approach:
- TensorFlow Universal Sentence Encoder (high-quality text embeddings) - 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 - Custom embedding functions can be plugged in for specialized domains
## Extensions ## Extensions
@ -1240,7 +1243,7 @@ const id = await db.addToBoth('Deep learning is a subset of machine learning', {
tags: ['deep learning', 'neural networks'] tags: ['deep learning', 'neural networks']
}) })
// Clean up when done // Clean up when done (this also cleans up worker pools)
await db.shutDown() 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: 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: 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 run build` (via the prebuild script)
- Running `npm version` commands (patch, minor, major) - Running `npm version` commands (patch, minor, major)
- Manually running `node scripts/generate-version.js` - 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. This ensures that the badge always reflects the current version in package.json, even before publishing to npm.

152
demo/index.html
View file

@ -1448,7 +1448,10 @@ await db.init()
await database.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() updateDatabaseVisualization()
// Populate noun dropdowns after initialization // Populate noun dropdowns after initialization
@ -1475,40 +1478,40 @@ await db.init()
type: 'person', type: 'person',
noun: 'person', noun: 'person',
name: 'Alex Johnson', name: 'Alex Johnson',
content: 'Tech enthusiast and software developer. Love hiking and photography. Working on AI projects in my spare time.', 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'] tags: ['tech', 'developer', 'photography', 'hiking', 'AI', 'machine learning']
}, },
{ {
id: 'user2', id: 'user2',
type: 'person', type: 'person',
noun: 'person', noun: 'person',
name: 'Sophia Chen', name: 'Sophia Chen',
content: 'Digital marketing specialist with a passion for data analytics. Foodie and travel blogger on weekends.', 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'] tags: ['marketing', 'analytics', 'food', 'travel', 'statistics', 'blogging']
}, },
{ {
id: 'user3', id: 'user3',
type: 'person', type: 'person',
noun: 'person', noun: 'person',
name: 'Marcus Williams', name: 'Marcus Williams',
content: 'Graphic designer and illustrator. Creating visual stories that connect people. Coffee addict.', 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'] tags: ['design', 'illustration', 'art', 'coffee', 'branding', 'minimalism']
}, },
{ {
id: 'user4', id: 'user4',
type: 'person', type: 'person',
noun: 'person', noun: 'person',
name: 'Priya Patel', name: 'Priya Patel',
content: 'Environmental scientist studying climate change impacts. Advocate for sustainable living and renewable energy.', 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'] tags: ['science', 'environment', 'sustainability', 'climate', 'research', 'education']
}, },
{ {
id: 'user5', id: 'user5',
type: 'person', type: 'person',
noun: 'person', noun: 'person',
name: 'Jordan Taylor', name: 'Jordan Taylor',
content: 'Fitness coach and nutrition expert. Helping people achieve their health goals through balanced lifestyle.', 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'] tags: ['fitness', 'nutrition', 'health', 'wellness', 'coaching', 'plant-based']
}, },
// Posts/Content // Posts/Content
@ -1517,40 +1520,40 @@ await db.init()
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'The Future of AI in Everyday Applications', 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.', 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'] tags: ['AI', 'technology', 'future', 'personalization', 'machine learning', 'privacy', 'ethics']
}, },
{ {
id: 'post2', id: 'post2',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'My Southeast Asia Food Tour', title: 'My Southeast Asia Food Tour: A Culinary Journey',
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.', 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'] tags: ['food', 'travel', 'asia', 'culinary', 'culture', 'cooking', 'Thailand', 'Vietnam', 'Malaysia']
}, },
{ {
id: 'post3', id: 'post3',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'Minimalist Design Principles for Digital Interfaces', title: 'Minimalist Design Principles for Digital Interfaces: Beyond Aesthetics',
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.', 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'] tags: ['design', 'minimalism', 'UI', 'UX', 'accessibility', 'typography', 'user experience', 'digital interfaces']
}, },
{ {
id: 'post4', id: 'post4',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'Urban Gardening: Growing Food in Small Spaces', 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.', 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'] tags: ['gardening', 'sustainability', 'urban', 'food', 'containers', 'vertical gardening', 'organic', 'small spaces']
}, },
{ {
id: 'post5', id: 'post5',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: '30-Day Home Workout Challenge', title: '30-Day Home Workout Challenge: Progressive Fitness Without Equipment',
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.', 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'] tags: ['fitness', 'workout', 'health', 'challenge', 'home exercise', 'bodyweight', 'functional fitness', 'progressive training']
}, },
// Comments // Comments
@ -1559,24 +1562,24 @@ await db.init()
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'Comment on AI post', 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?', 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'] tags: ['comment', 'AI', 'healthcare', 'privacy', 'federated learning', 'edge computing']
}, },
{ {
id: 'comment2', id: 'comment2',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'Comment on food tour', title: 'Comment on food tour',
content: 'Your photos are amazing! What was your favorite dish from the entire trip?', 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'] tags: ['comment', 'food', 'travel', 'photography', 'Vietnam', 'Malaysia', 'cooking', 'recipes']
}, },
{ {
id: 'comment3', id: 'comment3',
type: 'document', type: 'document',
noun: 'content', noun: 'content',
title: 'Comment on design article', title: 'Comment on design article',
content: 'Minimalism is definitely underrated. I\'ve found that my conversion rates improved significantly after simplifying our landing page.', 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'] tags: ['comment', 'design', 'conversion', 'business', 'accessibility', 'minimalism', 'typography', 'e-commerce']
}, },
// Groups/Communities // Groups/Communities
@ -1585,24 +1588,24 @@ await db.init()
type: 'group', type: 'group',
noun: 'group', noun: 'group',
name: 'Tech Innovators Network', name: 'Tech Innovators Network',
content: 'A community of developers, designers, and tech enthusiasts discussing emerging technologies and their applications.', 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'] tags: ['technology', 'innovation', 'community', 'networking', 'AI', 'blockchain', 'IoT', 'collaboration']
}, },
{ {
id: 'group2', id: 'group2',
type: 'group', type: 'group',
noun: 'group', noun: 'group',
name: 'Global Foodies', name: 'Global Foodies Collective',
content: 'Sharing culinary experiences, recipes, and food photography from around the world.', 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'] tags: ['food', 'cooking', 'international', 'recipes', 'culture', 'sustainability', 'traditions', 'culinary']
}, },
{ {
id: 'group3', id: 'group3',
type: 'group', type: 'group',
noun: 'group', noun: 'group',
name: 'Sustainable Living Collective', name: 'Sustainable Living Collective',
content: 'Exploring practical ways to reduce environmental impact and live more sustainably in modern society.', 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'] tags: ['sustainability', 'environment', 'lifestyle', 'eco-friendly', 'zero-waste', 'community', 'gardening', 'advocacy']
}, },
// Events // Events
@ -1610,28 +1613,29 @@ await db.init()
id: 'event1', id: 'event1',
type: 'event', type: 'event',
noun: 'event', noun: 'event',
name: 'Virtual Tech Meetup 2023', name: 'Virtual Tech Meetup 2023: Responsible AI Development',
content: 'Online gathering of tech professionals sharing insights on latest development frameworks and tools.', 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'] tags: ['tech', 'virtual', 'networking', 'learning', 'AI', 'ethics', 'development', 'workshops']
}, },
{ {
id: 'event2', id: 'event2',
type: 'event', type: 'event',
noun: 'event', noun: 'event',
name: 'International Food Festival', name: 'International Food Festival: Culinary Traditions in a Modern World',
content: 'Annual celebration of global cuisines featuring cooking demonstrations, tastings, and cultural performances.', 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'] tags: ['food', 'festival', 'international', 'culture', 'cooking', 'sustainability', 'traditions', 'innovation']
} }
] ]
// Add Nouns to the database // Add Nouns to the database using batch processing for better performance
for (const item of sampleNouns) { const nounItems = sampleNouns.map(item => ({
// Pass the content as the data to vectorize, and the entire item as metadata vectorOrData: item.content || item.name,
await database.add(item.content || item.name, item) metadata: item
} }))
// Wait a moment to ensure nouns are fully indexed before adding verbs await database.addBatch(nounItems, { batchSize: 50 })
await new Promise(resolve => setTimeout(resolve, 500))
// No need to wait as addBatch ensures all items are fully processed
// Verify nouns exist before adding verbs // Verify nouns exist before adding verbs
const entities = await database.getAllNouns() const entities = await database.getAllNouns()
@ -1921,24 +1925,38 @@ await db.init()
} }
] ]
// Add Verbs to the database // Add Verbs to the database using parallel processing for better performance
for (const item of sampleVerbs) { // First filter out verbs with invalid source or target nouns
try { const validVerbs = sampleVerbs.filter(item => {
// Check if source and target nouns exist before adding the verb if (!nounIds.includes(item.source)) {
if (!nounIds.includes(item.source)) { console.error(`Source noun with ID ${item.source} not found. Available nouns: ${nounIds.join(', ')}`)
throw new 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)) {
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}`)
} }
} 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 // Count successfully added verbs
const addedVerbs = await database.getAllVerbs() const addedVerbs = await database.getAllVerbs()

1
package.json
View file

@ -100,6 +100,7 @@
"dependencies": { "dependencies": {
"@aws-sdk/client-s3": "^3.427.0", "@aws-sdk/client-s3": "^3.427.0",
"@tensorflow-models/universal-sentence-encoder": "^1.3.3", "@tensorflow-models/universal-sentence-encoder": "^1.3.3",
"@tensorflow/tfjs": "^4.22.0",
"@tensorflow/tfjs-backend-cpu": "^4.22.0", "@tensorflow/tfjs-backend-cpu": "^4.22.0",
"@tensorflow/tfjs-backend-webgl": "^4.22.0", "@tensorflow/tfjs-backend-webgl": "^4.22.0",
"@tensorflow/tfjs-converter": "^4.22.0", "@tensorflow/tfjs-converter": "^4.22.0",

155
src/brainyData.ts
View file

@ -24,7 +24,9 @@ import {
import { import {
cosineDistance, cosineDistance,
defaultEmbeddingFunction, defaultEmbeddingFunction,
euclideanDistance defaultBatchEmbeddingFunction,
euclideanDistance,
cleanupWorkerPools
} from './utils/index.js' } from './utils/index.js'
import { NounType, VerbType, GraphNoun } from './types/graphTypes.js' import { NounType, VerbType, GraphNoun } from './types/graphTypes.js'
import { import {
@ -127,6 +129,7 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
private index: HNSWIndex | HNSWIndexOptimized private index: HNSWIndex | HNSWIndexOptimized
private storage: StorageAdapter | null = null private storage: StorageAdapter | null = null
private isInitialized = false private isInitialized = false
private isInitializing = false
private embeddingFunction: EmbeddingFunction private embeddingFunction: EmbeddingFunction
private distanceFunction: DistanceFunction private distanceFunction: DistanceFunction
private requestPersistentStorage: boolean private requestPersistentStorage: boolean
@ -204,7 +207,48 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
return return
} }
// Prevent recursive initialization
if (this.isInitializing) {
return
}
this.isInitializing = true
try { 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 // Initialize storage if not provided in constructor
if (!this.storage) { if (!this.storage) {
// Combine storage config with requestPersistentStorage for backward compatibility // Combine storage config with requestPersistentStorage for backward compatibility
@ -251,8 +295,10 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
} }
this.isInitialized = true this.isInitialized = true
this.isInitializing = false
} catch (error) { } catch (error) {
console.error('Failed to initialize BrainyData:', error) console.error('Failed to initialize BrainyData:', error)
this.isInitializing = false
throw new Error(`Failed to initialize BrainyData: ${error}`) throw new Error(`Failed to initialize BrainyData: ${error}`)
} }
} }
@ -478,6 +524,7 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
forceEmbed?: boolean // Force using the embedding function even if input is a vector 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 addToRemote?: boolean // Whether to also add to the remote server if connected
concurrency?: number // Maximum number of concurrent operations (default: 4) concurrency?: number // Maximum number of concurrent operations (default: 4)
batchSize?: number // Maximum number of items to process in a single batch (default: 50)
} = {} } = {}
): Promise<string[]> { ): Promise<string[]> {
await this.ensureInitialized() await this.ensureInitialized()
@ -488,22 +535,86 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
// Default concurrency to 4 if not specified // Default concurrency to 4 if not specified
const concurrency = options.concurrency || 4 const concurrency = options.concurrency || 4
// Default batch size to 50 if not specified
const batchSize = options.batchSize || 50
try { try {
// Process items in batches to control concurrency // Process items in batches to control concurrency and memory usage
const ids: string[] = [] const ids: string[] = []
const itemsToProcess = [...items] // Create a copy to avoid modifying the original array const itemsToProcess = [...items] // Create a copy to avoid modifying the original array
while (itemsToProcess.length > 0) { while (itemsToProcess.length > 0) {
// Take up to 'concurrency' items to process in parallel // Take up to 'batchSize' items to process in a batch
const batch = itemsToProcess.splice(0, concurrency) const batch = itemsToProcess.splice(0, batchSize)
// Process this batch in parallel // Separate items that are already vectors from those that need embedding
const batchResults = await Promise.all( const vectorItems: Array<{
batch.map((item) => vectorOrData: Vector
this.add(item.vectorOrData, item.metadata, options) 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<string>[] = []
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 // Add the results to our ids array
ids.push(...batchResults) ids.push(...batchResults)
} }
@ -1729,7 +1840,28 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
* Ensure the database is initialized * Ensure the database is initialized
*/ */
private async ensureInitialized(): Promise<void> { private async ensureInitialized(): Promise<void> {
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() await this.init()
} }
} }
@ -1838,6 +1970,9 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
await this.disconnectFromRemoteServer() await this.disconnectFromRemoteServer()
} }
// Clean up worker pools to release resources
cleanupWorkerPools()
// Additional cleanup could be added here in the future // Additional cleanup could be added here in the future
this.isInitialized = false this.isInitialized = false

15
src/global.d.ts vendored
View file

@ -5,11 +5,16 @@
// Extend the globalThis interface to include the __ENV__ property // Extend the globalThis interface to include the __ENV__ property
declare global { declare global {
var __ENV__: { var __ENV__: {
isBrowser: boolean; isBrowser: boolean
isNode: string | false; isNode: string | false
isServerless: boolean; isServerless: boolean
}; }
// Window interface extensions (worker-specific functions removed)
interface Window {
importTensorFlow?: () => Promise<any>
}
} }
// This export is needed to make this file a module // This export is needed to make this file a module
export {}; export {}

52
src/hnsw/hnswIndex.ts
View file

@ -10,7 +10,7 @@ import {
Vector, Vector,
VectorDocument VectorDocument
} from '../coreTypes.js' } from '../coreTypes.js'
import { euclideanDistance, calculateDistancesWithGPU } from '../utils/index.js' import { euclideanDistance, calculateDistancesBatch } from '../utils/index.js'
import { executeInThread } from '../utils/workerUtils.js' import { executeInThread } from '../utils/workerUtils.js'
// Default HNSW parameters // Default HNSW parameters
@ -83,7 +83,7 @@ export class HNSWIndex {
const vectorsOnly = vectors.map((item) => item.vector) const vectorsOnly = vectors.map((item) => item.vector)
// Use GPU-accelerated distance calculation when possible // Use GPU-accelerated distance calculation when possible
const distances = await calculateDistancesWithGPU( const distances = await calculateDistancesBatch(
queryVector, queryVector,
vectorsOnly, vectorsOnly,
this.distanceFunction this.distanceFunction
@ -96,51 +96,15 @@ export class HNSWIndex {
})) }))
} catch (error) { } catch (error) {
console.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 error
) )
// Fall back to threaded CPU processing if GPU acceleration fails // Fall back to sequential processing if GPU acceleration fails
// Function to be executed in a worker thread return vectors.map((item) => ({
const distanceCalculator = (args: { id: item.id,
queryVector: Vector distance: this.distanceFunction(queryVector, item.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<Array<{ id: string; distance: number }>>(
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)
}))
}
} }
} }

310
src/utils/distance.ts
View file

@ -13,7 +13,10 @@ import { isThreadingAvailable } from './environment.js'
* Lower values indicate higher similarity * Lower values indicate higher similarity
* Optimized using array methods for Node.js 23.11+ * 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) { if (a.length !== b.length) {
throw new Error('Vectors must have the same dimensions') 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+ // Use array.reduce for better performance in Node.js 23.11+
const sum = a.reduce((acc, val, i) => { const sum = a.reduce((acc, val, i) => {
const diff = val - b[i] const diff = val - b[i]
return acc + (diff * diff) return acc + diff * diff
}, 0) }, 0)
return Math.sqrt(sum) return Math.sqrt(sum)
@ -33,19 +36,25 @@ export const euclideanDistance: DistanceFunction = (a: Vector, b: Vector): numbe
* Range: 0 (identical) to 2 (opposite) * Range: 0 (identical) to 2 (opposite)
* Optimized using array methods for Node.js 23.11+ * 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) { if (a.length !== b.length) {
throw new Error('Vectors must have the same dimensions') throw new Error('Vectors must have the same dimensions')
} }
// Use array.reduce to calculate all values in a single pass // Use array.reduce to calculate all values in a single pass
const { dotProduct, normA, normB } = a.reduce((acc, val, i) => { const { dotProduct, normA, normB } = a.reduce(
return { (acc, val, i) => {
dotProduct: acc.dotProduct + (val * b[i]), return {
normA: acc.normA + (val * val), dotProduct: acc.dotProduct + val * b[i],
normB: acc.normB + (b[i] * b[i]) normA: acc.normA + val * val,
} normB: acc.normB + b[i] * b[i]
}, { dotProduct: 0, normA: 0, normB: 0 }) }
},
{ dotProduct: 0, normA: 0, normB: 0 }
)
if (normA === 0 || normB === 0) { if (normA === 0 || normB === 0) {
return 2 // Maximum distance for zero vectors 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 * Lower values indicate higher similarity
* Optimized using array methods for Node.js 23.11+ * 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) { if (a.length !== b.length) {
throw new Error('Vectors must have the same dimensions') 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) * Converted to a distance metric (lower is better)
* Optimized using array methods for Node.js 23.11+ * 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) { if (a.length !== b.length) {
throw new Error('Vectors must have the same dimensions') throw new Error('Vectors must have the same dimensions')
} }
// Use array.reduce for better performance in Node.js 23.11+ // 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) // Convert to a distance metric (lower is better)
return -dotProduct return -dotProduct
} }
/** /**
* GPU-accelerated batch distance calculation * Batch distance calculation
* Uses TensorFlow.js with WebGL backend when available for optimal performance * Uses TensorFlow.js with CPU backend for optimized performance
* Falls back to CPU processing when GPU is not available
* *
* @param queryVector The query vector to compare against all vectors * @param queryVector The query vector to compare against all vectors
* @param vectors Array of vectors to compare against * @param vectors Array of vectors to compare against
* @param distanceFunction The distance function to use * @param distanceFunction The distance function to use
* @returns Promise resolving to array of distances * @returns Promise resolving to array of distances
*/ */
export async function calculateDistancesWithGPU( export async function calculateDistancesBatch(
queryVector: Vector, queryVector: Vector,
vectors: Vector[], vectors: Vector[],
distanceFunction: DistanceFunction = euclideanDistance distanceFunction: DistanceFunction = euclideanDistance
): Promise<number[]> { ): Promise<number[]> {
// For small batches, use the standard distance function // For small batches, use the standard distance function
if (vectors.length < 10) { if (vectors.length < 10) {
return vectors.map(vector => distanceFunction(queryVector, vector)) return vectors.map((vector) => distanceFunction(queryVector, vector))
} }
try { try {
// Function to be executed in a worker thread // Function to be executed in a worker thread
const distanceCalculator = async ( const distanceCalculator = async (args: {
args: { queryVector: Vector
queryVector: Vector, vectors: Vector[]
vectors: Vector[], distanceFnString: string
distanceFnString: string }) => {
}
) => {
const { queryVector, vectors, distanceFnString } = args const { queryVector, vectors, distanceFnString } = args
// Try to use TensorFlow.js with GPU acceleration if available // Use TensorFlow.js with CPU processing
const useTensorFlow = async () => { 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 // 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') await import('@tensorflow/tfjs-backend-cpu')
let usingGPU = false // Set CPU as the backend
await tf.setBackend('cpu')
}
try { // Convert vectors to tensors
// Try to import and use WebGL backend (GPU) const queryTensor = tf.tensor2d([queryVector])
await import('@tensorflow/tfjs-backend-webgl') const vectorsTensor = tf.tensor2d(vectors)
// Check if WebGL is available and set it as the backend let distances: number[]
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')
}
// Convert vectors to tensors // Calculate distances based on the distance function type
const queryTensor = tf.tensor2d([queryVector]) if (distanceFnString.includes('euclideanDistance')) {
const vectorsTensor = tf.tensor2d(vectors) // 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 const queryNorm = tf.norm(queryTensor, 2, 1)
if (distanceFnString.includes('euclideanDistance')) { const vectorsNorm = tf.norm(vectorsTensor, 2, 1)
// 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[]
// Clean up tensors const normProduct = tf.outerProduct(
queryTensor.dispose() queryNorm as any,
vectorsTensor.dispose() vectorsNorm as any
expanded.dispose() )
squaredDiff.dispose() const cosineSimilarity = tf.div(dotProduct, normProduct)
sumSquaredDiff.dispose() const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity)
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 queryNorm = tf.norm(queryTensor, 2, 1) distances = (await (distancesTensor as any)
const vectorsNorm = tf.norm(vectorsTensor, 2, 1) .squeeze()
.array()) as number[]
const normProduct = tf.outerProduct(queryNorm as any, vectorsNorm as any) // Clean up tensors
const cosineSimilarity = tf.div(dotProduct, normProduct) queryTensor.dispose()
const distancesTensor = tf.sub(tf.scalar(1), cosineSimilarity) 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 // Clean up tensors
queryTensor.dispose() queryTensor.dispose()
vectorsTensor.dispose() vectorsTensor.dispose()
dotProduct.dispose() diff.dispose()
queryNorm.dispose() absDiff.dispose()
vectorsNorm.dispose() distancesTensor.dispose()
normProduct.dispose() } else if (distanceFnString.includes('dotProductDistance')) {
cosineSimilarity.dispose() // Dot product distance using GPU-optimized operations
distancesTensor.dispose() // Formula: -sum(a * b)
} else if (distanceFnString.includes('manhattanDistance')) { const dotProduct = tf.matMul(
// Manhattan distance using GPU-optimized operations queryTensor,
// Formula: sum(|a - b|) (vectorsTensor as any).transpose()
const diff = tf.sub((queryTensor as any).expandDims(1), (vectorsTensor as any).expandDims(0)) )
const absDiff = tf.abs(diff) const distancesTensor = tf.neg(dotProduct)
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 // Clean up tensors
queryTensor.dispose() queryTensor.dispose()
vectorsTensor.dispose() vectorsTensor.dispose()
diff.dispose() dotProduct.dispose()
absDiff.dispose() distancesTensor.dispose()
distancesTensor.dispose() } else {
} else if (distanceFnString.includes('dotProductDistance')) { // For unknown distance functions, fall back to direct CPU implementation
// Dot product distance using GPU-optimized operations throw new Error(
// Formula: -sum(a * b) 'Unsupported distance function for TensorFlow optimization'
const dotProduct = tf.matMul(queryTensor, (vectorsTensor as any).transpose()) )
const distancesTensor = tf.neg(dotProduct) }
distances = await (distancesTensor as any).squeeze().array() as number[] return {
distances
// 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
} }
} }
// Try to use TensorFlow.js with GPU acceleration // Try to use TensorFlow.js with CPU optimization
try { try {
return await useTensorFlow() return await useTensorFlow()
} catch (error) { } 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 // 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 // Calculate distances for all vectors
const distances = vectors.map(vector => distanceFunction(queryVector, vector)) const distances = vectors.map((vector) =>
distanceFunction(queryVector, vector)
)
return { return {
distances, distances
usingGPU: false
} }
} }
} }
// Execute the distance calculation in a separate thread if threading is available // Threading is not available, so we'll always use the main thread implementation
if (isThreadingAvailable()) { // This comment is kept for clarity about the removed code
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)
}
}
// If threading is not available or failed, calculate distances in the main thread // 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) { } catch (error) {
// If anything fails, fall back to the standard distance function // If anything fails, fall back to the standard distance function
console.error('GPU-accelerated distance calculation failed:', error) console.error('Batch distance calculation failed:', error)
return vectors.map(vector => distanceFunction(queryVector, vector)) return vectors.map((vector) => distanceFunction(queryVector, vector))
} }
} }

452
src/utils/embedding.ts
View file

@ -10,14 +10,15 @@ import { executeInThread } from './workerUtils.js'
* This model provides high-quality text embeddings using TensorFlow.js * This model provides high-quality text embeddings using TensorFlow.js
* The required TensorFlow.js dependencies are automatically installed with this package * The required TensorFlow.js dependencies are automatically installed with this package
* *
* This implementation will use GPU acceleration via WebGL when available, * This implementation attempts to use GPU processing when available for better performance,
* falling back to CPU processing when GPU is not available or fails to initialize. * falling back to CPU processing for compatibility across all environments.
*/ */
export class UniversalSentenceEncoder implements EmbeddingModel { export class UniversalSentenceEncoder implements EmbeddingModel {
private model: any = null private model: any = null
private initialized = false private initialized = false
private tf: any = null private tf: any = null
private use: any = null private use: any = null
private backend: string = 'cpu' // Default to CPU
/** /**
* Initialize the embedding model * Initialize the embedding model
@ -47,36 +48,59 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// Use type assertions to tell TypeScript these modules exist // Use type assertions to tell TypeScript these modules exist
this.tf = await import('@tensorflow/tfjs-core') 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') await import('@tensorflow/tfjs-backend-cpu')
// Try to import WebGL backend for GPU acceleration in browser environments
try { try {
// Try to import and use WebGL backend first (GPU) if (typeof window !== 'undefined') {
await import('@tensorflow/tfjs-backend-webgl') await import('@tensorflow/tfjs-backend-webgl')
// Check if WebGL is available using setBackend instead of findBackend
// Check if WebGL is available and set it as the backend try {
if ( if (this.tf.setBackend) {
(await this.tf.findBackend('webgl')) || await this.tf.setBackend('webgl')
(await this.tf.ready().then(() => this.tf.findBackend('webgl'))) this.backend = 'webgl'
) { console.log('Using WebGL backend for TensorFlow.js')
console.log('Using WebGL backend (GPU acceleration)') } else {
await this.tf.setBackend('webgl') console.warn(
} else { 'tf.setBackend is not available, falling back to CPU'
console.log('WebGL backend not available, falling back to CPU') )
await this.tf.setBackend('cpu') }
} catch (e) {
console.warn('WebGL backend not available, falling back to CPU:', e)
this.backend = 'cpu'
}
} }
} catch (err) { } catch (error) {
console.warn( console.warn('WebGL backend not available, falling back to CPU:', error)
'WebGL backend failed to initialize, using CPU backend:', this.backend = 'cpu'
err }
)
await this.tf.setBackend('cpu') // Set the backend
if (this.tf.setBackend) {
await this.tf.setBackend(this.backend)
} }
this.use = await import('@tensorflow-models/universal-sentence-encoder') 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 // Load the model
this.model = await this.use.load() this.model = await loadFunction()
this.initialized = true this.initialized = true
// Restore original console.warn // Restore original console.warn
@ -132,6 +156,10 @@ export class UniversalSentenceEncoder implements EmbeddingModel {
// Convert to array and return the first embedding // Convert to array and return the first embedding
const embeddingArray = await embeddings.array() const embeddingArray = await embeddings.array()
// Dispose of the tensor to free memory
embeddings.dispose()
return embeddingArray[0] return embeddingArray[0]
} catch (error) { } catch (error) {
console.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<Vector[]> {
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 * 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 * Create an embedding function from an embedding model
* @param model Embedding model to use * @param model Embedding model to use
@ -177,21 +368,22 @@ export function createEmbeddingFunction(
/** /**
* Creates a TensorFlow-based Universal Sentence Encoder embedding function * Creates a TensorFlow-based Universal Sentence Encoder embedding function
* This is the required embedding function for all text embeddings * 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 that persists across all embedding calls
// Create a single shared instance of the model const sharedModel = new UniversalSentenceEncoder()
const model = new UniversalSentenceEncoder() let sharedModelInitialized = false
let modelInitialized = false
export function createTensorFlowEmbeddingFunction(): EmbeddingFunction {
return async (data: any): Promise<Vector> => { return async (data: any): Promise<Vector> => {
try { try {
// Initialize the model if it hasn't been initialized yet // Initialize the model if it hasn't been initialized yet
if (!modelInitialized) { if (!sharedModelInitialized) {
await model.init() await sharedModel.init()
modelInitialized = true sharedModelInitialized = true
} }
return await model.embed(data) return await sharedModel.embed(data)
} catch (error) { } catch (error) {
console.error('Failed to use TensorFlow embedding:', error) console.error('Failed to use TensorFlow embedding:', error)
throw new 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 * Default embedding function
* This provides better performance for embedding operations by: * Uses UniversalSentenceEncoder for all text embeddings
* 1. Using GPU acceleration via WebGL when available * TensorFlow.js is required for this to work
* 2. Running in a separate thread to avoid blocking the main thread * Uses CPU for compatibility
* 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
*/ */
export function createThreadedEmbeddingFunction( export const defaultEmbeddingFunction: EmbeddingFunction =
options: { fallbackToMain?: boolean } = {} createTensorFlowEmbeddingFunction()
): EmbeddingFunction {
// Create a standard embedding function to use as fallback
const standardEmbedding = 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<Vector> => { export const defaultBatchEmbeddingFunction: (
// If we've already determined that threading doesn't work, use the fallback dataArray: string[]
if (useFallback) { ) => Promise<Vector[]> = async (dataArray: string[]): Promise<Vector[]> => {
return standardEmbedding(data) try {
// Initialize the model if it hasn't been initialized yet
if (!sharedBatchModelInitialized) {
await sharedBatchModel.init()
sharedBatchModelInitialized = true
} }
try { return await sharedBatchModel.embedBatch(dataArray)
// Function to be executed in a worker thread } catch (error) {
// This must be a regular function (not async) to avoid Promise cloning issues console.error('Failed to use TensorFlow batch embedding:', error)
const embedInWorker = (inputData: any) => { throw new Error(
// Return a plain object with the input data `Universal Sentence Encoder batch embedding failed: ${error}`
// 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<Vector>(
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}`)
}
} }
} }
/** /**
* Default embedding function * Creates an embedding function that runs in a separate thread
* Uses UniversalSentenceEncoder for all text embeddings * This is a wrapper around createEmbeddingFunction that uses executeInThread
* TensorFlow.js is required for this to work * @param model Embedding model to use
* 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
*/ */
export const defaultEmbeddingFunction: EmbeddingFunction = export function createThreadedEmbeddingFunction(
createThreadedEmbeddingFunction({ fallbackToMain: true }) model: EmbeddingModel
): EmbeddingFunction {
const embeddingFunction = createEmbeddingFunction(model)
return async (data: any): Promise<Vector> => {
// 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<Vector>(fnString, data)
}
}

5
src/utils/environment.ts
View file

@ -39,7 +39,7 @@ export function areWebWorkersAvailable(): boolean {
*/ */
export function areWorkerThreadsAvailable(): boolean { export function areWorkerThreadsAvailable(): boolean {
if (!isNode()) return false; if (!isNode()) return false;
try { try {
// Dynamic import to avoid errors in browser environments // Dynamic import to avoid errors in browser environments
require('worker_threads'); require('worker_threads');
@ -51,7 +51,8 @@ export function areWorkerThreadsAvailable(): boolean {
/** /**
* Determine if threading is available in the current environment * Determine if threading is available in the current environment
* Always returns false since multithreading has been removed
*/ */
export function isThreadingAvailable(): boolean { export function isThreadingAvailable(): boolean {
return areWebWorkersAvailable() || areWorkerThreadsAvailable(); return false;
} }

1
src/utils/index.ts
View file

@ -1,2 +1,3 @@
export * from './distance.js' export * from './distance.js'
export * from './embedding.js' export * from './embedding.js'
export * from './workerUtils.js'

162
src/utils/workerUtils.ts
View file

@ -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) * 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 executeInWebWorker<T>(fnString: string, args: any): Promise<T> {
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<T>(
fnString: string,
args: any
): Promise<T> {
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
* *
* @param fnString The function to execute as a string * @param fnString The function to execute as a string
* @param args The arguments to pass to the function * @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function * @returns A promise that resolves with the result of the function
*/ */
export function executeInThread<T>(fnString: string, args: any): Promise<T> { export function executeInThread<T>(fnString: string, args: any): Promise<T> {
if (isBrowser()) { try {
return executeInWebWorker<T>(fnString, args)
} else if (isNode()) {
return executeInWorkerThread<T>(fnString, args)
} else {
// Fall back to executing in the main thread
// Parse the function from string and execute it // Parse the function from string and execute it
const fn = new Function('return ' + fnString)() const fn = new Function('return ' + fnString)()
return Promise.resolve(fn(args) as T) 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)')
}