├── .github └── workflows │ └── build.yml ├── .gitignore ├── .npmignore ├── LICENSE ├── README.md ├── benchmark.ts ├── index.ts ├── package.json ├── test.ts └── tsconfig.json /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | build: 7 | runs-on: >- 8 | ${{ 9 | (matrix.os == 'mac' && matrix.arch == 'arm64') && 10 | 'macos-14' || 11 | (fromJson('{"linux":"ubuntu-22.04","mac":"macos-13","win":"windows-2022"}')[matrix.os]) 12 | }} 13 | continue-on-error: false 14 | 15 | strategy: 16 | fail-fast: false 17 | matrix: 18 | os: [linux, mac] 19 | arch: [x64] 20 | include: 21 | - os: mac 22 | arch: arm64 23 | 24 | steps: 25 | - name: Setup Node.js 26 | uses: actions/setup-node@v4 27 | with: 28 | node-version: 22 29 | 30 | - name: Install mac dependencies 31 | if: matrix.os == 'mac' && matrix.arch == 'x64' 32 | run: brew install openblas 33 | 34 | - name: Checkout 35 | uses: actions/checkout@v4 36 | 37 | - name: Test 38 | run: | 39 | yarn 40 | yarn test 41 | 42 | publish: 43 | if: startsWith(github.ref, 'refs/tags/') 44 | needs: [build] 45 | runs-on: ubuntu-latest 46 | 47 | steps: 48 | - name: Checkout 49 | uses: actions/checkout@v4 50 | with: 51 | submodules: true 52 | 53 | - name: Get tag 54 | run: echo "VERSION=${GITHUB_REF#refs/*/}" >> $GITHUB_ENV 55 | 56 | - name: Set package version 57 | run: | 58 | npm config set git-tag-version=false 59 | npm version $VERSION 60 | 61 | - name: Install deps 62 | run: yarn 63 | 64 | - name: Publish npm package 65 | uses: JS-DevTools/npm-publish@v3 66 | with: 67 | token: ${{ secrets.NPM_TOKEN }} 68 | access: public 69 | ignore-scripts: false 70 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # TypeScript compiled files 2 | /dist/ 3 | 4 | # Everything else should keep same with .npmignore 5 | *.swp 6 | *.tgz 7 | 8 | yarn.lock 9 | package-lock.json 10 | npm-debug.log 11 | yarn-error.log 12 | /node_modules/ 13 | -------------------------------------------------------------------------------- /.npmignore: -------------------------------------------------------------------------------- 1 | # Unused source files 2 | /.github/ 3 | /dist/benchmark.* 4 | /dist/test.* 5 | 6 | # Everything else should keep same with .gitignore 7 | *.swp 8 | *.tgz 9 | 10 | yarn.lock 11 | package-lock.json 12 | npm-debug.log 13 | yarn-error.log 14 | /node_modules/ 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright © 2024 zcbenz 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # You need a few lines of JS, not a vector database 2 | 3 | Suppose you are implementing RAG for your AI app, and you have used web APIs or 4 | an inference engine to generate a large number of embeddings, and now you need 5 | to find out the best results matching a query embedding, what do you do? 6 | 7 | Use a vector database? No, you only need a few lines of JavaScript. 8 | 9 | ## A quick introduction to node-mlx 10 | 11 | [MLX](https://github.com/ml-explore/mlx) is a full-featured machine learning 12 | framework, with easy-to-understand source code and 13 | [small binary sizes](https://github.com/frost-beta/node-mlx/releases). 14 | And [node-mlx](https://github.com/frost-beta/node-mlx) is the JavaScript binding 15 | of it. 16 | 17 | MLX only has GPU support for macOS, and but its CPU support, implemented with 18 | vectorized instructions, is still fast on Linux. 19 | 20 | ```typescript 21 | import {core as mx, nn} from '@frost-beta/mlx'; 22 | ``` 23 | 24 | ## Embeddings search, in a few lines of JS 25 | 26 | Suppose you want to find out the results with highest similaries to `query`, 27 | from the `embeddings`. 28 | 29 | ```typescript 30 | const embeddings = [ 31 | [ 0.037035, 0.0760545, ... ], 32 | [ 0.034029, -0.0227216, ... ], 33 | ... 34 | [ -0.028612, 0.0052857, ... ], 35 | ]; 36 | const query = [ -0.019773, 0.006021, ... ]; 37 | ``` 38 | 39 | With node-mlx, you can use the builtin `nn.losses.cosineSimilarityLoss` API to 40 | do the search. 41 | 42 | ```typescript 43 | const embeddingsTensor = mx.array(embeddings); 44 | const queryTensor = mx.array([ query ]); 45 | const scoresTensor = nn.losses.cosineSimilarityLoss(queryTensor, embeddingsTensor); 46 | const scores: Float32Array = scoresTensor.toTypedArray(); 47 | ``` 48 | 49 | The `scores` array stores the 50 | [cosine similarities](https://en.wikipedia.org/wiki/Cosine_similarity) 51 | between the `query` and `embeddings`. 52 | 53 | (If you are wondering how we can compute cosine similarities between a 1x1 54 | tensor and a 1xN tensor, it is called 55 | [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html).) 56 | 57 | ```typescript 58 | console.log(scores); 59 | // [ 0.233244, 0.012492, ..., 0.43298 ] 60 | ``` 61 | 62 | ## Sorting 63 | 64 | Once you get the `scores` array, you can use the usual JavScript code to filter 65 | and sort the results. But you can also use MLX if the number of results is large 66 | enough to make JavScript engine struggle. 67 | 68 | ```typescript 69 | // Get the top 10 scores. 70 | const topTen = mx.topk(scoresTensor, 10); 71 | console.log(topTen.toTypedArray()); 72 | // [ 0.894323, 0.594729, ... ] 73 | 74 | // Sort the scores. 75 | let sortedScores = mx.sort(scoresTensor); 76 | console.log(sortedScores.toTypedArray()); 77 | // [ 0.01287, 0.1502876, ... ] 78 | sortedScores = sortedScores.index(mx.Slice(null, null, -1)); 79 | console.log(sortedScores.toTypedArray()); 80 | // [ 0.894323, 0.594729, ... ] 81 | 82 | // Get the indices of the scores ordered by their values in the array. 83 | const indices = mx.argsort(scoresTensor) 84 | .index(mx.Slice(null, null, -1)) 85 | .toTypedArray(); 86 | console.log(indices); 87 | // [ 8, 9, ... ] 88 | console.log(indices.map(i => scores[i])); 89 | // [ 0.894323, 0.594729, ... ] 90 | ``` 91 | 92 | The `array.index(mx.Slice(null, null, -1))` code looks alien, it is actually the 93 | JavaScript version of Python's `array[::-1]`, which reverse the array. You can 94 | of course convert the result to JavaScript Array frist and then call 95 | `reverse()`, but it would be slower if the array is very large. 96 | 97 | ## A Node.js module 98 | 99 | If after reading above introductions you still find MXL cumbersome to use (which 100 | is normal if you had zero experience with NumPy or PyTorch), I have wrapped the 101 | code into a very simple Node.js module, which you can use to replace vector 102 | databases in many cases. 103 | 104 | Install: 105 | 106 | ```console 107 | npm install not-a-vector-database 108 | ``` 109 | 110 | APIs: 111 | 112 | ```typescript 113 | export type EmbeddingInput = number[] | TypedArray; 114 | 115 | export interface SearchOptions { 116 | /** 117 | * Results with scores larger than this value will be returned. Should be 118 | * between -1 and 1. Default is 0. 119 | */ 120 | minimumScore?: number; 121 | /** 122 | * Restrict the number of results to return, default is 16. 123 | */ 124 | maximumResults?: number; 125 | } 126 | 127 | export interface SearchResult { 128 | score: number; 129 | data: unknown; 130 | } 131 | 132 | /** 133 | * In-memory storage of embeddings and associated data. 134 | */ 135 | export class Storage { 136 | embeddings?: mx.array; 137 | data: unknown[]; 138 | 139 | /** 140 | * Initialize from the buffer. 141 | */ 142 | loadFromBuffer(buffer: Buffer): void; 143 | 144 | /** 145 | * Dump the data to a buffer. 146 | */ 147 | dumpToBuffer(): Buffer; 148 | 149 | /** 150 | * Add data to the storage. 151 | */ 152 | push(...items: {embedding: EmbeddingInput; data: unknown;}[]): void; 153 | 154 | /** 155 | * Return the data which are most relevant to the embedding. 156 | */ 157 | search(embedding: EmbeddingInput, options?: SearchOptions): SearchResult[]; 158 | } 159 | ``` 160 | 161 | Example: 162 | 163 | ```typescript 164 | import {Storage} from 'not-a-vector-database'; 165 | 166 | const storage = new Stroage(); 167 | storage.push({embedding, data: 'some data'}); 168 | 169 | const results = storage.search(embedding); 170 | 171 | fs.writeFileSync('storage.bser', storage.dumpToBuffer()); 172 | storage.loadFromBuffer(fs.readFileSync('storage.bser')); 173 | ``` 174 | 175 | There is also a `benchmark.ts` script that you can use to test the performance. 176 | On my 2018 Intel MacBook Pro which has no GPU support, searching from 1 million 177 | embeddings with size of 128 takes about 900ms. 178 | -------------------------------------------------------------------------------- /benchmark.ts: -------------------------------------------------------------------------------- 1 | import {core as mx} from '@frost-beta/mlx'; 2 | import {Storage} from './index.js'; 3 | 4 | const embeddingSize = 128; 5 | const totalDataSize = 1000 * 1000; 6 | 7 | console.log('Preparing 1M embeddings...'); 8 | const embeddings = mx.random.uniform(-1, 1, [ totalDataSize, embeddingSize ]); 9 | const query = mx.random.uniform(-1, 1, [ embeddingSize ]); 10 | mx.eval(embeddings, query); 11 | 12 | const storage = new Storage(); 13 | storage.embeddings = embeddings; 14 | storage.data = new Array(totalDataSize); 15 | 16 | console.log('Searching...'); 17 | const input = query.toTypedArray(); 18 | console.time('Time'); 19 | const results = storage.search(input); 20 | console.timeEnd('Time'); 21 | -------------------------------------------------------------------------------- /index.ts: -------------------------------------------------------------------------------- 1 | import {core as mx, nn} from '@frost-beta/mlx'; 2 | import * as bser from 'bser'; 3 | 4 | type TypedArray = Int8Array | Uint8Array | Int16Array | Uint16Array | 5 | Int32Array | Uint32Array | Float32Array | Float64Array; 6 | export type EmbeddingInput = number[] | TypedArray | mx.array; 7 | 8 | export interface SearchOptions { 9 | /** 10 | * Results with scores larger than this value will be returned. Should be 11 | * between -1 and 1. Default is 0. 12 | */ 13 | minimumScore?: number; 14 | /** 15 | * Restrict the number of results to return, default is 16. 16 | */ 17 | maximumResults?: number; 18 | } 19 | 20 | export interface SearchResult { 21 | score: number; 22 | data: unknown; 23 | } 24 | 25 | /** 26 | * In-memory storage of embeddings and associated data. 27 | */ 28 | export class Storage { 29 | embeddings?: mx.array; 30 | data: unknown[] = []; 31 | 32 | /** 33 | * Initialize from the buffer. 34 | */ 35 | loadFromBuffer(buffer: Buffer) { 36 | const json = bser.loadFromBuffer(buffer) as any; 37 | if (!Array.isArray(json.embeddings) && !Array.isArray(json.data)) 38 | throw new Error('The buffer does not includes valid data.'); 39 | if (this.embeddings) 40 | mx.dispose(this.embeddings); 41 | this.embeddings = mx.array(json.embeddings); 42 | this.data = json.data; 43 | } 44 | 45 | /** 46 | * Dump the data to a buffer. 47 | */ 48 | dumpToBuffer(): Buffer { 49 | if (!this.embeddings) 50 | throw new Error('There is no data in storage.'); 51 | return bser.dumpToBuffer({ 52 | embeddings: this.embeddings.tolist(), 53 | data: this.data, 54 | }); 55 | } 56 | 57 | /** 58 | * Add data to the storage. 59 | */ 60 | push(...items: {embedding: EmbeddingInput, data: unknown}[]) { 61 | // Make sure intermediate arrays are released. 62 | mx.tidy(() => { 63 | const newEmbeddings = mx.stack(items.map(i => mx.array(i.embedding))); 64 | if (!this.embeddings) { 65 | this.embeddings = newEmbeddings; 66 | } else { 67 | const old = this.embeddings; 68 | this.embeddings = mx.concatenate([ old, newEmbeddings ], 0); 69 | // Release the old embeddings object, which is not caught by tidy. 70 | mx.dispose(old); 71 | } 72 | // Do not release this.embeddings. 73 | return this.embeddings; 74 | }); 75 | this.data.push(...items.map(i => i.data)); 76 | } 77 | 78 | /** 79 | * Return the data which are most relevant to the embedding. 80 | */ 81 | search(embedding: EmbeddingInput, 82 | { 83 | minimumScore = 0, 84 | maximumResults = 16, 85 | }: SearchOptions = {}): SearchResult[] { 86 | if (!this.embeddings) 87 | return []; 88 | return mx.tidy(() => { 89 | const query = mx.array(embedding, this.embeddings.dtype).index(mx.newaxis); 90 | const scores = nn.losses.cosineSimilarityLoss(query, this.embeddings); 91 | const indices = mx.argsort(scores).index(mx.Slice(null, null, -1)) 92 | .index(mx.Slice(null, maximumResults)) 93 | .toTypedArray(); 94 | const results: SearchResult[] = []; 95 | for (const index of indices) { 96 | const score = scores.index(index).item() as number; 97 | if (score < minimumScore) 98 | break; 99 | results.push({score, data: this.data[index]}); 100 | } 101 | return results; 102 | }); 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "not-a-vector-database", 3 | "version": "0.0.1-dev", 4 | "description": "Store and search embeddings", 5 | "main": "dist/index.js", 6 | "types": "dist/index.d.ts", 7 | "scripts": { 8 | "prepack": "tsc", 9 | "pretest": "tsc", 10 | "test": "node --test dist/test.js" 11 | }, 12 | "author": "zcbenz", 13 | "license": "MIT", 14 | "keywords": [ "mlx", "embeddings", "vector database" ], 15 | "repository": { 16 | "type": "git", 17 | "url": "git+https://github.com/frost-beta/not-a-vector-database.git" 18 | }, 19 | "bugs": { 20 | "url": "https://github.com/frost-beta/not-a-vector-database/issues" 21 | }, 22 | "devDependencies": { 23 | "@types/bser": "2.0.4", 24 | "@types/node": "22.5.4", 25 | "typescript": "5.6.2" 26 | }, 27 | "dependencies": { 28 | "@frost-beta/mlx": "0.0.21", 29 | "bser": "2.1.1" 30 | } 31 | } 32 | -------------------------------------------------------------------------------- /test.ts: -------------------------------------------------------------------------------- 1 | import {describe, it} from 'node:test'; 2 | import * as assert from 'node:assert'; 3 | import {core as mx} from '@frost-beta/mlx'; 4 | 5 | import {Storage} from './index.js'; 6 | 7 | function generateEmbedding() { 8 | const embeddingSize = 16; 9 | return mx.random.uniform(-1, 1, [ embeddingSize ]).toTypedArray() as Float32Array; 10 | } 11 | 12 | describe('Storage', () => { 13 | it('push', () => { 14 | const storage = new Storage(); 15 | storage.push({embedding: generateEmbedding(), data: 1}); 16 | assert.deepEqual(storage.embeddings.shape, [ 1, 16 ]); 17 | storage.push({embedding: generateEmbedding(), data: 2}, 18 | {embedding: generateEmbedding(), data: 3}); 19 | assert.deepEqual(storage.embeddings.shape, [ 3, 16 ]); 20 | assert.deepEqual(storage.data, [ 1, 2, 3 ]); 21 | }); 22 | 23 | it('search', () => { 24 | const size = 10; 25 | const storage = new Storage(); 26 | for (let i = 0; i < size; ++i) { 27 | storage.push({embedding: generateEmbedding(), data: `data${i}`}); 28 | } 29 | const results = storage.search(generateEmbedding()); 30 | const scores = results.map(r => r.score); 31 | assert.ok(scores.every(s => s >= 0)); 32 | assert.deepEqual(scores, scores.toSorted().toReversed()); 33 | }); 34 | 35 | it('serialization', () => { 36 | const size = 10; 37 | const storage = new Storage(); 38 | const embeddings: Float32Array[] = []; 39 | const data: string[] = []; 40 | for (let i = 0; i < size; ++i) { 41 | embeddings.push(generateEmbedding()); 42 | data.push(`data${i}`); 43 | } 44 | for (let i = 0; i < size; ++i) { 45 | storage.push({embedding: embeddings[i], data: data[i]}); 46 | } 47 | storage.loadFromBuffer(storage.dumpToBuffer()); 48 | for (let i = 0; i < size; ++i) { 49 | assert.deepEqual(storage.embeddings.index(i).tolist(), Array.from(embeddings[i])); 50 | assert.equal(storage.data[i], data[i]); 51 | } 52 | }); 53 | }); 54 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "outDir": "dist", 4 | "declaration": true, 5 | "allowJs": false, 6 | "module": "commonjs", 7 | "resolveJsonModule": true, 8 | "target": "es2023", 9 | "lib": [ "esnext" ] 10 | }, 11 | "include": [ "*.ts" ], 12 | "exclude": [ "node_modules" ] 13 | } 14 | --------------------------------------------------------------------------------