Merge pull request #70 from Cloud-Code-AI/53-add-embedding-support

feat: Added embedding models
Cloud-Code-AI · Feb 2, 2025 · 21a7084 · 21a7084
2 parents 3c770d4 + e0ea227
commit 21a7084
Show file tree

Hide file tree

Showing 8 changed files with 146 additions and 13 deletions.
diff --git a/README.md b/README.md
@@ -141,6 +141,10 @@ More models will be added soon. Request a model by creating an issue.
 - Qwen2.5-1.5B-Instruct
 - DeepSeek-R1-Distill-Qwen-7B
 - DeepSeek-R1-Distill-Llama-8B
+- Snowflake-Arctic-Embed-M-B32
+- Snowflake-Arctic-Embed-S-B32
+- Snowflake-Arctic-Embed-M-B4
+- Snowflake-Arctic-Embed-S-B4
 
 ### Transformers Models
 - Llama-3.2-1b-Instruct

diff --git a/examples/database-demos/indexeddb-demo/src/components/DatabaseDemo.tsx b/examples/database-demos/indexeddb-demo/src/components/DatabaseDemo.tsx
@@ -2,13 +2,15 @@ import { useEffect, useState } from 'react';
 import { DatabaseImpl } from '@browserai/browserai';
 import { Document } from '../types';
 import './DatabaseDemo.css';
+import { BrowserAI } from '@browserai/browserai';
 
 export function DatabaseDemo() {
     const [logs, setLogs] = useState<string[]>([]);
     const [db, setDb] = useState<DatabaseImpl<Document> | null>(null);
     const [inputText, setInputText] = useState('');
     const [documentId, setDocumentId] = useState('');
     const [documents, setDocuments] = useState<Document[]>([]);
+    const [embeddings, setEmbeddings] = useState<BrowserAI | null>(null);
 
     const addLog = (message: string) => {
         setLogs(prev => [...prev, message]);
@@ -17,28 +19,38 @@ export function DatabaseDemo() {
     useEffect(() => {
         const initDb = async () => {
             const config = {
-                databaseName: 'myAppDB',
-                version: 1
+                databaseName: 'BrowserAI-DB',
+                version: 1,
+                stores: {
+                    documents: { keyPath: 'id' }
+                }
             };
             try {
                 const database = new DatabaseImpl<Document>({ type: 'indexeddb', config });
-                // Wait for initialization to complete
                 await database.initialize({ type: 'indexeddb', config });
+
+                const browserAI = new BrowserAI();
+                await browserAI.loadModel('snowflake-arctic-embed-m-b32');
+
                 setDb(database);
+                setEmbeddings(browserAI);
+                addLog('✅ Database and embeddings initialized');
             } catch (error) {
-                addLog(`❌ Error initializing database: ${error instanceof Error ? error.message : String(error)}`);
+                addLog(`❌ Error initializing: ${error instanceof Error ? 'An error occurred during initialization.' : 'An unknown error occurred.'}`);
             }
         };
 
         initDb();
 
         return () => {
-            // Only call close if db is initialized
             if (db) {
                 db.close();
             }
+            if (embeddings) {
+                addLog('✅ Embeddings closed');
+            }
         };
-    }, []); // Remove db from dependencies to avoid recreation
+    }, []);
 
     const refreshDocuments = async () => {
         if (!db) return;
@@ -47,10 +59,17 @@ export function DatabaseDemo() {
     };
 
     const handleStore = async () => {
-        if (!db || !inputText || !documentId) return;
+        if (!db || !embeddings || !inputText || !documentId) return;
         try {
-            await db.store({ id: documentId, text: inputText });
-            addLog(`✅ Stored document with ID: ${documentId}`);
+            const embedding = await embeddings.embed(inputText);
+            console.log(embedding);
+            await db.store({ 
+                id: documentId, 
+                text: inputText,
+                embedding: embedding
+            });
+
+            addLog(`✅ Stored document with ID: ${documentId} and embedding`);
             await refreshDocuments();
             setInputText('');
             setDocumentId('');
@@ -80,10 +99,17 @@ export function DatabaseDemo() {
     };
 
     const handleUpdate = async () => {
-        if (!db || !inputText || !documentId) return;
+        if (!db || !embeddings || !inputText || !documentId) return;
         try {
-            await db.update({ id: documentId, text: inputText });
-            addLog(`🔄 Updated document with ID: ${documentId}`);
+            const embedding = await embeddings.embed(inputText);
+
+            await db.update({ 
+                id: documentId, 
+                text: inputText,
+                embedding: embedding
+            });
+
+            addLog(`🔄 Updated document with ID: ${documentId} and new embedding`);
             await refreshDocuments();
             setInputText('');
             setDocumentId('');

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@browserai/browserai",
-  "version": "1.0.20",
+  "version": "1.0.21",
   "private": false,
   "description": "A library for running AI models directly in the browser",
   "main": "dist/index.js",

diff --git a/src/config/models/mlc-models.json b/src/config/models/mlc-models.json
@@ -176,5 +176,85 @@
         ],
         "defaultQuantization": "q4f16_1",
         "pipeline": "text-generation"
+    },
+    "snowflake-arctic-embed-m-b4": {
+        "engine": "mlc",
+        "modelName": "snowflake-arctic-embed-m-b4",
+        "modelType": "embedding",
+        "repo": "mlc-ai/snowflake-arctic-embed-m-{quantization}-MLC-b4",
+        "quantizations": [
+            "q0f32"
+        ],
+        "defaultQuantization": "q0f32",
+        "pipeline": "embedding",
+        "metadata": {
+            "context_window_size": 512,
+            "download_size_in_mb": {
+                "q0f32": 218
+            },
+            "estimated_vram_in_mb": {
+                "q0f32": 218
+            }
+        }
+    },
+    "snowflake-arctic-embed-s-b4": {
+        "engine": "mlc",
+        "modelName": "snowflake-arctic-embed-s-b4",
+        "modelType": "embedding",
+        "repo": "mlc-ai/snowflake-arctic-embed-s-{quantization}-MLC-b4",
+        "quantizations": [
+            "q0f32"
+        ],
+        "defaultQuantization": "q0f32",
+        "pipeline": "embedding",
+        "metadata": {
+            "context_window_size": 512,
+            "download_size_in_mb": {
+                "q0f32": 66.4
+            },
+            "estimated_vram_in_mb": {
+                "q0f32": 66.4
+            }
+        }
+    },
+    "snowflake-arctic-embed-m-b32": {
+        "engine": "mlc",
+        "modelName": "snowflake-arctic-embed-m-b32",
+        "modelType": "embedding",
+        "repo": "mlc-ai/snowflake-arctic-embed-m-{quantization}-MLC-b32",
+        "quantizations": [
+            "q0f32"
+        ],
+        "defaultQuantization": "q0f32",
+        "pipeline": "embedding",
+        "metadata": {
+            "context_window_size": 512,
+            "download_size_in_mb": {
+                "q0f32": 218
+            },
+            "estimated_vram_in_mb": {
+                "q0f32": 218
+            }
+        }
+    },
+    "snowflake-arctic-embed-s-b32": {
+        "engine": "mlc",
+        "modelName": "snowflake-arctic-embed-s-b32",
+        "modelType": "embedding",
+        "repo": "mlc-ai/snowflake-arctic-embed-s-{quantization}-MLC-b32",
+        "quantizations": [
+            "q0f32"
+        ],
+        "defaultQuantization": "q0f32",
+        "pipeline": "embedding",
+        "metadata": {
+            "context_window_size": 512,
+            "download_size_in_mb": {
+                "q0f32": 66.4
+            },
+            "estimated_vram_in_mb": {
+                "q0f32": 66.4
+            }
+        }
     }
 }
diff --git a/src/config/models/types.ts b/src/config/models/types.ts
@@ -10,6 +10,7 @@ export interface BaseModelConfig {
   quantizations?: string[];
   requiredFeatures?: string[];
   modelLibrary?: string;
+  metadata?: Record<string, any>;
 }
 
 export type ModelType =

diff --git a/src/core/llm/index.ts b/src/core/llm/index.ts
@@ -73,6 +73,13 @@ export class BrowserAI {
     }
   }
 
+  async embed(input: string, options: Record<string, unknown> = {}): Promise<unknown> {
+    if (!this.engine) {
+      throw new Error('No model loaded. Please call loadModel first.');
+    }
+    return await this.engine.embed(input, options);
+  }
+
   async transcribeAudio(audio: Blob | Float32Array, options: Record<string, unknown> = {}): Promise<unknown> {
     if (!this.engine) {
       throw new Error('No model loaded. Please call loadModel first.');

diff --git a/src/engines/mlc-engine-wrapper.ts b/src/engines/mlc-engine-wrapper.ts
@@ -70,4 +70,12 @@ export class MLCEngineWrapper {
     const result = await this.mlcEngine.chat.completions.create({ messages, ...options });
     return result.choices[0].message.content;
   }
+
+  async embed(input: string, options: any = {}) {
+    if (!this.mlcEngine) {
+      throw new Error('MLC Engine not initialized.');
+    }
+    const result = await this.mlcEngine.embeddings.create({ input, ...options });
+    return result.data[0].embedding;
+  }
 }
diff --git a/src/engines/transformer-engine-wrapper.ts b/src/engines/transformer-engine-wrapper.ts
@@ -161,4 +161,11 @@ export class TransformersEngineWrapper {
         throw new Error(`Unsupported model type: ${this.modelType}`);
     }
   }
+
+  async embed(input: string, options: any = {}) {
+    if (!this.transformersPipeline || this.modelType !== 'feature-extraction') {
+      console.debug(`Feature extraction pipeline not initialized. ${input}, ${options}`);
+      throw new Error('Feature extraction pipeline not initialized.');
+    }
+  }
 }