diff --git a/README.md b/README.md index 69451a63..c3feda1c 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ A vector database that runs in a browser or Node.js and utilizes Origin Private - **Persistent storage**: Uses Origin Private File System (OPFS) in browsers, with fallback to in-memory storage - **Efficient vector search**: Implements HNSW (Hierarchical Navigable Small World) algorithm for fast approximate nearest neighbor search - **Automatic embedding**: Converts text and other data to vectors using embedding models -- **TensorFlow.js integration**: Uses Universal Sentence Encoder for high-quality text embeddings +- **TensorFlow.js integration**: Uses Universal Sentence Encoder for high-quality text embeddings (TensorFlow.js is included as a dependency) - **Metadata support**: Store and retrieve metadata alongside vectors - **TypeScript support**: Fully typed API with generics for metadata types - **Multiple distance functions**: Supports cosine, Euclidean, Manhattan, and dot product distance metrics @@ -95,11 +95,41 @@ console.log(catVector); const animalVectors = await db.embed(["cat", "dog", "fish"]); console.log(animalVectors); // [0.123, 0.456, 0.789, ...] - Vector representation of the first item in the array +``` -// You can also use the exported embedding classes directly -import {UniversalSentenceEncoder, createEmbeddingFunction} from '@soulcraft/brainy'; +### Using Embedding Functions -// Create a new Universal Sentence Encoder model +By default, Brainy uses the TensorFlow Universal Sentence Encoder for high-quality text embeddings. The TensorFlow.js dependencies are automatically included when you install the package, so you don't need to install them separately. + +You can use the default embedding function directly: + +```typescript +import { + defaultEmbeddingFunction, + createTensorFlowEmbeddingFunction, + createSimpleEmbeddingFunction, + UniversalSentenceEncoder, + createEmbeddingFunction +} from '@soulcraft/brainy'; + +// Option 1: Use the default embedding function (TensorFlow Universal Sentence Encoder) +const vector1 = await defaultEmbeddingFunction("Some text to embed"); +console.log(vector1); +// [0.123, 0.456, 0.789, ...] - High-quality vector representation using TensorFlow + +// Option 2: Explicitly create a TensorFlow-based embedding function +const tfEmbedFunction = createTensorFlowEmbeddingFunction(); +const vector2 = await tfEmbedFunction("Some text to embed"); +console.log(vector2); +// [0.123, 0.456, 0.789, ...] - High-quality vector representation using TensorFlow + +// Option 3: Use the simple character-based embedding (faster but less accurate) +const simpleEmbedFunction = createSimpleEmbeddingFunction(); +const vector3 = await simpleEmbedFunction("Some text to embed"); +console.log(vector3); +// [0.123, 0.456, 0.789, ...] - Basic vector representation using character frequencies + +// Option 4: Create the model and embedding function manually const useModel = new UniversalSentenceEncoder(); await useModel.init(); @@ -107,14 +137,27 @@ await useModel.init(); const embedFunction = createEmbeddingFunction(useModel); // Embed text using the function -const vector = await embedFunction("Some text to embed"); -console.log(vector); -// [0.123, 0.456, 0.789, ...] - Vector representation of "Some text to embed" +const vector4 = await embedFunction("Some text to embed"); +console.log(vector4); +// [0.123, 0.456, 0.789, ...] - High-quality vector representation using TensorFlow // Don't forget to dispose of the model when done await useModel.dispose(); ``` +You can also configure BrainyData to use a different embedding function if needed: + +```typescript +import { BrainyData, createSimpleEmbeddingFunction } from '@soulcraft/brainy'; + +// Create a new vector database with the simple embedding function +// (only if you prefer speed over accuracy) +const db = new BrainyData({ + embeddingFunction: createSimpleEmbeddingFunction() +}); +await db.init(); +``` + ### Configuration Options ```typescript @@ -724,7 +767,9 @@ constructor(config?: BrainyDataConfig) ### Embedding Functions - `createEmbeddingFunction(model: EmbeddingModel): EmbeddingFunction` - Create an embedding function from an embedding model -- `defaultEmbeddingFunction` - Default embedding function using UniversalSentenceEncoder +- `createTensorFlowEmbeddingFunction(): EmbeddingFunction` - Create an embedding function using TensorFlow's Universal Sentence Encoder +- `createSimpleEmbeddingFunction(): EmbeddingFunction` - Create a simple character-based embedding function (faster but less accurate) +- `defaultEmbeddingFunction` - Default embedding function using TensorFlow's Universal Sentence Encoder for high-quality embeddings ## Browser Compatibility diff --git a/package.json b/package.json index a7760a5f..e3b4fa82 100644 --- a/package.json +++ b/package.json @@ -50,12 +50,12 @@ "typescript": "^5.1.6" }, "dependencies": { + "uuid": "^9.0.0", "@tensorflow-models/universal-sentence-encoder": "^1.3.3", "@tensorflow/tfjs": "^4.22.0", "@tensorflow/tfjs-backend-cpu": "^4.22.0", "@tensorflow/tfjs-core": "^4.22.0", - "@tensorflow/tfjs-layers": "^4.22.0", - "uuid": "^9.0.0" + "@tensorflow/tfjs-layers": "^4.22.0" }, "prettier": { "arrowParens": "always", diff --git a/src/index.ts b/src/index.ts index 5e51e679..8fd10203 100644 --- a/src/index.ts +++ b/src/index.ts @@ -27,12 +27,16 @@ import { SimpleEmbedding, UniversalSentenceEncoder, createEmbeddingFunction, + createTensorFlowEmbeddingFunction, + createSimpleEmbeddingFunction, defaultEmbeddingFunction } from './utils/embedding.js' export { SimpleEmbedding, UniversalSentenceEncoder, createEmbeddingFunction, + createTensorFlowEmbeddingFunction, + createSimpleEmbeddingFunction, defaultEmbeddingFunction } diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index 0d145c77..6e101003 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -75,7 +75,8 @@ export class SimpleEmbedding implements EmbeddingModel { /** * TensorFlow Universal Sentence Encoder embedding model - * Requires @tensorflow/tfjs and @tensorflow-models/universal-sentence-encoder to be installed + * This model provides high-quality text embeddings using TensorFlow.js + * The required TensorFlow.js dependencies are automatically installed with this package */ export class UniversalSentenceEncoder implements EmbeddingModel { private model: any = null @@ -98,7 +99,9 @@ export class UniversalSentenceEncoder implements EmbeddingModel { this.initialized = true } catch (error) { console.error('Failed to initialize Universal Sentence Encoder:', error) - throw new Error(`Failed to initialize Universal Sentence Encoder: ${error}`) + throw new Error( + `Failed to initialize Universal Sentence Encoder: ${error}` + ) } } @@ -163,6 +166,23 @@ export function createEmbeddingFunction(model: EmbeddingModel): EmbeddingFunctio } /** - * Default embedding function using UniversalSentenceEncoder + * Creates a TensorFlow-based Universal Sentence Encoder embedding function + * This is the recommended embedding function for high-quality text embeddings */ -export const defaultEmbeddingFunction: EmbeddingFunction = createEmbeddingFunction(new UniversalSentenceEncoder()) +export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { + return createEmbeddingFunction(new UniversalSentenceEncoder()) +} + +/** + * Simple embedding function using character-based embedding + * This is a basic implementation that doesn't use TensorFlow + */ +export function createSimpleEmbeddingFunction(): EmbeddingFunction { + return createEmbeddingFunction(new SimpleEmbedding()) +} + +/** + * Default embedding function using UniversalSentenceEncoder + * This provides high-quality text embeddings using TensorFlow.js + */ +export const defaultEmbeddingFunction: EmbeddingFunction = createTensorFlowEmbeddingFunction()