├── .cargo └── config.toml ├── .gitattributes ├── .github └── workflows │ └── rust.yml ├── .gitignore ├── .vscode └── launch.json ├── Cargo.toml ├── LICENSE.md ├── README.md ├── _rust-toolchain.toml ├── precommit.sh └── src ├── database ├── core.rs ├── default │ ├── audio.rs │ ├── image.rs │ ├── mod.rs │ └── text.rs ├── index │ ├── lsh.rs │ └── mod.rs └── mod.rs ├── distance.rs ├── lib.rs ├── main.rs └── model ├── audio.rs ├── core.rs ├── image.rs ├── mod.rs └── text.rs /.cargo/config.toml: -------------------------------------------------------------------------------- 1 | [doc.extern-map.registries] 2 | crates-io = "https://docs.rs/" 3 | 4 | [target.aarch64-apple-darwin] 5 | rustflags = ["-C", "target-cpu=native", "-C", "strip=symbols"] 6 | 7 | [target.x86_64-apple-darwin] 8 | rustflags = ["-C", "target-cpu=native", "-C", "strip=symbols"] 9 | 10 | [target.x86_64-pc-windows-gnu] 11 | rustflags = ["-C", "link-arg=-lpsapi", "-C", "link-arg=-lbcrypt", "-C", "target-cpu=native", "-C", "strip=symbols"] 12 | 13 | [target.aarch64-unknown-linux-gnu] 14 | rustflags = ["-C", "target-cpu=native", "-C", "strip=symbols"] 15 | 16 | [target.x86_64-unknown-linux-gnu] 17 | rustflags = ["-C", "target-cpu=native", "-C", "strip=symbols"] -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.github/workflows/rust.yml: -------------------------------------------------------------------------------- 1 | name: Rust 2 | on: 3 | push: 4 | branches: [ "master" ] 5 | env: 6 | CARGO_TERM_COLOR: always 7 | jobs: 8 | build_documentation: 9 | name: Build documentation 10 | runs-on: ubuntu-latest 11 | steps: 12 | - name: Setup Rust toolchain 13 | uses: actions-rs/toolchain@v1 14 | with: 15 | toolchain: nightly 16 | target: x86_64-unknown-linux-gnu 17 | default: true 18 | profile: default 19 | - name: Checkout codebase 20 | uses: actions/checkout@v4 21 | - name: Generate documentation 22 | run: time cargo doc --features="default_db" --no-deps -Zrustdoc-map --release --quiet 23 | - name: Fix permissions 24 | run: | 25 | chmod -c -R +rX "target/doc/" | while read line; do 26 | echo "::warning title=Invalid file permissions automatically fixed::$line" 27 | done 28 | - name: Upload Pages artifact 29 | uses: actions/upload-pages-artifact@v3 30 | with: 31 | path: "target/doc/" 32 | deploy_documentation: 33 | needs: build_documentation 34 | name: Deploy documentation to GitHub Pages 35 | permissions: 36 | pages: write 37 | id-token: write 38 | environment: 39 | name: github-pages 40 | url: ${{ steps.deployment.outputs.page_url }} 41 | runs-on: ubuntu-latest 42 | steps: 43 | - name: Deploy to GitHub Pages 44 | id: deployment 45 | uses: actions/deploy-pages@v4 46 | apply_suggestions: 47 | name: Format code, apply compiler suggestions 48 | runs-on: ubuntu-latest 49 | steps: 50 | - name: Checkout codebase 51 | uses: actions/checkout@v4 52 | - name: Setup Rust toolchain 53 | uses: actions-rs/toolchain@v1 54 | with: 55 | toolchain: nightly 56 | components: clippy, rustfmt 57 | profile: minimal 58 | - name: Format 59 | run: cargo fmt 60 | - name: Apply compiler suggestions 61 | run: | 62 | cargo clippy --fix --allow-dirty --features="cli" 63 | - name: Commit changes to code, if any 64 | run: | 65 | git config user.name github-actions 66 | git config user.email github-actions@github.com 67 | git diff --quiet && git diff --staged --quiet || git commit -am "chore: Format and apply compiler suggestions." 68 | git push 69 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ## Linux 2 | *~ 3 | 4 | # temporary files which can be created if a process still has a handle open of a deleted file 5 | .fuse_hidden* 6 | 7 | # KDE directory preferences 8 | .directory 9 | 10 | # Linux trash folder which might appear on any partition or disk 11 | .Trash-* 12 | 13 | # .nfs files are created when an open file is removed but is still being accessed 14 | .nfs* 15 | 16 | ## macOS 17 | # General 18 | .DS_Store 19 | .AppleDouble 20 | .LSOverride 21 | 22 | # Icon must end with two \r 23 | Icon 24 | 25 | 26 | # Thumbnails 27 | ._* 28 | 29 | # Files that might appear in the root of a volume 30 | .DocumentRevisions-V100 31 | .fseventsd 32 | .Spotlight-V100 33 | .TemporaryItems 34 | .Trashes 35 | .VolumeIcon.icns 36 | .com.apple.timemachine.donotpresent 37 | 38 | # Directories potentially created on remote AFP share 39 | .AppleDB 40 | .AppleDesktop 41 | Network Trash Folder 42 | Temporary Items 43 | .apdisk 44 | 45 | ## Windows 46 | # Windows thumbnail cache files 47 | Thumbs.db 48 | Thumbs.db:encryptable 49 | ehthumbs.db 50 | ehthumbs_vista.db 51 | 52 | # Dump file 53 | *.stackdump 54 | 55 | # Folder config file 56 | [Dd]esktop.ini 57 | 58 | # Recycle Bin used on file shares 59 | $RECYCLE.BIN/ 60 | 61 | # Windows Installer files 62 | *.cab 63 | *.msi 64 | *.msix 65 | *.msm 66 | *.msp 67 | 68 | # Windows shortcuts 69 | *.lnk 70 | 71 | ## Visual Studio Code 72 | .vscode/* 73 | !.vscode/settings.json 74 | !.vscode/tasks.json 75 | !.vscode/launch.json 76 | !.vscode/extensions.json 77 | !.vscode/*.code-snippets 78 | 79 | # Local History for Visual Studio Code 80 | .history/ 81 | 82 | # Built Visual Studio Code Extensions 83 | *.vsix 84 | 85 | ## Rust 86 | # Generated by Cargo 87 | # will have compiled files and executables 88 | debug/ 89 | target/ 90 | 91 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 92 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 93 | Cargo.lock 94 | 95 | # These are backup files generated by rustfmt 96 | **/*.rs.bk 97 | 98 | # MSVC Windows builds of rustc generate these, which store debugging information 99 | *.pdb 100 | 101 | ## JetBrains 102 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider 103 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 104 | 105 | # User-specific stuff 106 | .idea/**/workspace.xml 107 | .idea/**/tasks.xml 108 | .idea/**/usage.statistics.xml 109 | .idea/**/dictionaries 110 | .idea/**/shelf 111 | 112 | # AWS User-specific 113 | .idea/**/aws.xml 114 | 115 | # Generated files 116 | .idea/**/contentModel.xml 117 | 118 | # Sensitive or high-churn files 119 | .idea/**/dataSources/ 120 | .idea/**/dataSources.ids 121 | .idea/**/dataSources.local.xml 122 | .idea/**/sqlDataSources.xml 123 | .idea/**/dynamic.xml 124 | .idea/**/uiDesigner.xml 125 | .idea/**/dbnavigator.xml 126 | 127 | # Gradle 128 | .idea/**/gradle.xml 129 | .idea/**/libraries 130 | 131 | # Gradle and Maven with auto-import 132 | # When using Gradle or Maven with auto-import, you should exclude module files, 133 | # since they will be recreated, and may cause churn. Uncomment if using 134 | # auto-import. 135 | # .idea/artifacts 136 | # .idea/compiler.xml 137 | # .idea/jarRepositories.xml 138 | # .idea/modules.xml 139 | # .idea/*.iml 140 | # .idea/modules 141 | # *.iml 142 | # *.ipr 143 | 144 | # CMake 145 | cmake-build-*/ 146 | 147 | # Mongo Explorer plugin 148 | .idea/**/mongoSettings.xml 149 | 150 | # File-based project format 151 | *.iws 152 | 153 | # IntelliJ 154 | out/ 155 | 156 | # mpeltonen/sbt-idea plugin 157 | .idea_modules/ 158 | 159 | # JIRA plugin 160 | atlassian-ide-plugin.xml 161 | 162 | # Cursive Clojure plugin 163 | .idea/replstate.xml 164 | 165 | # SonarLint plugin 166 | .idea/sonarlint/ 167 | 168 | # Crashlytics plugin (for Android Studio and IntelliJ) 169 | com_crashlytics_export_strings.xml 170 | crashlytics.properties 171 | crashlytics-build.properties 172 | fabric.properties 173 | 174 | # Editor-based Rest Client 175 | .idea/httpRequests 176 | 177 | # Android studio 3.1+ serialized cache file 178 | .idea/caches/build_file_checksums.ser 179 | 180 | ## Project-specific Ignores 181 | .fastembed_cache/ 182 | text.db 183 | texts 184 | image.db 185 | images 186 | audio.db 187 | audio 188 | *_old.rs 189 | .idea 190 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | { 2 | // Use IntelliSense to learn about possible attributes. 3 | // Hover to view descriptions of existing attributes. 4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 5 | "version": "0.2.0", 6 | "configurations": [ 7 | { 8 | "type": "lldb", 9 | "request": "launch", 10 | "name": "Debug unit tests in library 'zebra'", 11 | "cargo": { 12 | "args": [ 13 | "test", 14 | "--no-run", 15 | "--lib", 16 | "--package=zebra" 17 | ], 18 | "filter": { 19 | "name": "zebra", 20 | "kind": "lib" 21 | } 22 | }, 23 | "args": [], 24 | "cwd": "${workspaceFolder}" 25 | }, 26 | { 27 | "type": "lldb", 28 | "request": "launch", 29 | "name": "Debug executable 'zebra'", 30 | "cargo": { 31 | "args": [ 32 | "build", 33 | "--bin=zebra", 34 | "--package=zebra" 35 | ], 36 | "filter": { 37 | "name": "zebra", 38 | "kind": "bin" 39 | } 40 | }, 41 | "args": [], 42 | "cwd": "${workspaceFolder}" 43 | }, 44 | { 45 | "name": "(lldb) Attach", 46 | "type": "lldb", 47 | "request": "attach", 48 | "program": "${workspaceFolder}/target/release/zebra", 49 | }, 50 | { 51 | "type": "lldb", 52 | "request": "launch", 53 | "name": "Debug unit tests in executable 'zebra'", 54 | "cargo": { 55 | "args": [ 56 | "test", 57 | "--no-run", 58 | "--bin=zebra", 59 | "--package=zebra" 60 | ], 61 | "filter": { 62 | "name": "zebra", 63 | "kind": "bin" 64 | } 65 | }, 66 | "args": [], 67 | "cwd": "${workspaceFolder}" 68 | } 69 | ] 70 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "zebra" 3 | authors = ["Emil Sayahi"] 4 | description = "A vector database for querying meaningfully similar data." 5 | repository = "https://github.com/emmyoh/zebra/" 6 | license = "AGPL-3.0-or-later" 7 | readme = "README.md" 8 | version = "0.1.0" 9 | edition = "2021" 10 | 11 | [lib] 12 | name = "zebra" 13 | path = "src/lib.rs" 14 | crate-type = ["rlib", "dylib", "staticlib"] 15 | 16 | [[bin]] 17 | name = "zebra" 18 | path = "src/main.rs" 19 | doc = false 20 | required-features = ["cli"] 21 | 22 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 23 | 24 | [dependencies] 25 | clap = { version = "4.5.23", features = ["derive"], optional = true } 26 | fastembed = "4.3.0" 27 | simsimd = "6.2.3" 28 | space = "0.18.0" 29 | lz4_flex = { version = "0.11.3", default-features = false, features = [ 30 | "frame", 31 | ] } 32 | ticky = { version = "1.0.2", optional = true } 33 | pretty-duration = { version = "0.1.1", optional = true } 34 | indicatif = { version = "0.17.9", optional = true, features = [ 35 | "rayon", 36 | "improved_unicode", 37 | ] } 38 | distances = "1.8.0" 39 | candle-examples = "0.8.1" 40 | candle-core = "0.8.1" 41 | candle-nn = "0.8.1" 42 | candle-transformers = "0.8.1" 43 | hf-hub = "0.4.1" 44 | viuer = { version = "0.9.1", optional = true } 45 | sonogram = "0.7.1" 46 | image = "0.25.5" 47 | rodio = { version = "0.20.1", optional = true } 48 | rayon = "1.10.0" 49 | bytes = { version = "1.9.0" } 50 | symphonia = "0.5.4" 51 | anyhow = "1.0.95" 52 | hamming-bitwise-fast = "1.0.0" 53 | dashmap = { version = "6.1.0", features = ["rayon", "inline", "serde"] } 54 | rand = "0.9.0" 55 | fjall = "2.4.4" 56 | uuid = { version = "1.11.0", features = ["fast-rng", "v7", "serde"] } 57 | serde = { version = "1.0.217", features = ["derive"] } 58 | bincode = { version = "2.0.0", features = ["serde"] } 59 | serde_with = "3.12.0" 60 | 61 | [features] 62 | default = [] 63 | default_db = [] 64 | accelerate = [ 65 | "candle-core/accelerate", 66 | "candle-examples/accelerate", 67 | "candle-nn/accelerate", 68 | "candle-transformers/accelerate", 69 | ] 70 | cuda = [ 71 | "candle-core/cuda", 72 | "candle-examples/cuda", 73 | "candle-nn/cuda", 74 | "candle-transformers/cuda", 75 | ] 76 | mkl = [ 77 | "candle-core/mkl", 78 | "candle-examples/mkl", 79 | "candle-nn/mkl", 80 | "candle-transformers/mkl", 81 | ] 82 | metal = [ 83 | "candle-core/metal", 84 | "candle-examples/metal", 85 | "candle-nn/metal", 86 | "candle-transformers/metal", 87 | ] 88 | sixel = ["viuer/sixel"] 89 | cli = [ 90 | "default_db", 91 | "dep:clap", 92 | "dep:ticky", 93 | "dep:pretty-duration", 94 | "dep:indicatif", 95 | "dep:viuer", 96 | "dep:rodio", 97 | ] 98 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | # GNU AFFERO GENERAL PUBLIC LICENSE 2 | 3 | Version 3, 19 November 2007 4 | 5 | Copyright (C) 2007 Free Software Foundation, Inc. 6 | 7 | 8 | Everyone is permitted to copy and distribute verbatim copies of this 9 | license document, but changing it is not allowed. 10 | 11 | ## Preamble 12 | 13 | The GNU Affero General Public License is a free, copyleft license for 14 | software and other kinds of works, specifically designed to ensure 15 | cooperation with the community in the case of network server software. 16 | 17 | The licenses for most software and other practical works are designed 18 | to take away your freedom to share and change the works. By contrast, 19 | our General Public Licenses are intended to guarantee your freedom to 20 | share and change all versions of a program--to make sure it remains 21 | free software for all its users. 22 | 23 | When we speak of free software, we are referring to freedom, not 24 | price. Our General Public Licenses are designed to make sure that you 25 | have the freedom to distribute copies of free software (and charge for 26 | them if you wish), that you receive source code or can get it if you 27 | want it, that you can change the software or use pieces of it in new 28 | free programs, and that you know you can do these things. 29 | 30 | Developers that use our General Public Licenses protect your rights 31 | with two steps: (1) assert copyright on the software, and (2) offer 32 | you this License which gives you legal permission to copy, distribute 33 | and/or modify the software. 34 | 35 | A secondary benefit of defending all users' freedom is that 36 | improvements made in alternate versions of the program, if they 37 | receive widespread use, become available for other developers to 38 | incorporate. Many developers of free software are heartened and 39 | encouraged by the resulting cooperation. However, in the case of 40 | software used on network servers, this result may fail to come about. 41 | The GNU General Public License permits making a modified version and 42 | letting the public access it on a server without ever releasing its 43 | source code to the public. 44 | 45 | The GNU Affero General Public License is designed specifically to 46 | ensure that, in such cases, the modified source code becomes available 47 | to the community. It requires the operator of a network server to 48 | provide the source code of the modified version running there to the 49 | users of that server. Therefore, public use of a modified version, on 50 | a publicly accessible server, gives the public access to the source 51 | code of the modified version. 52 | 53 | An older license, called the Affero General Public License and 54 | published by Affero, was designed to accomplish similar goals. This is 55 | a different license, not a version of the Affero GPL, but Affero has 56 | released a new version of the Affero GPL which permits relicensing 57 | under this license. 58 | 59 | The precise terms and conditions for copying, distribution and 60 | modification follow. 61 | 62 | ## TERMS AND CONDITIONS 63 | 64 | ### 0. Definitions. 65 | 66 | "This License" refers to version 3 of the GNU Affero General Public 67 | License. 68 | 69 | "Copyright" also means copyright-like laws that apply to other kinds 70 | of works, such as semiconductor masks. 71 | 72 | "The Program" refers to any copyrightable work licensed under this 73 | License. Each licensee is addressed as "you". "Licensees" and 74 | "recipients" may be individuals or organizations. 75 | 76 | To "modify" a work means to copy from or adapt all or part of the work 77 | in a fashion requiring copyright permission, other than the making of 78 | an exact copy. The resulting work is called a "modified version" of 79 | the earlier work or a work "based on" the earlier work. 80 | 81 | A "covered work" means either the unmodified Program or a work based 82 | on the Program. 83 | 84 | To "propagate" a work means to do anything with it that, without 85 | permission, would make you directly or secondarily liable for 86 | infringement under applicable copyright law, except executing it on a 87 | computer or modifying a private copy. Propagation includes copying, 88 | distribution (with or without modification), making available to the 89 | public, and in some countries other activities as well. 90 | 91 | To "convey" a work means any kind of propagation that enables other 92 | parties to make or receive copies. Mere interaction with a user 93 | through a computer network, with no transfer of a copy, is not 94 | conveying. 95 | 96 | An interactive user interface displays "Appropriate Legal Notices" to 97 | the extent that it includes a convenient and prominently visible 98 | feature that (1) displays an appropriate copyright notice, and (2) 99 | tells the user that there is no warranty for the work (except to the 100 | extent that warranties are provided), that licensees may convey the 101 | work under this License, and how to view a copy of this License. If 102 | the interface presents a list of user commands or options, such as a 103 | menu, a prominent item in the list meets this criterion. 104 | 105 | ### 1. Source Code. 106 | 107 | The "source code" for a work means the preferred form of the work for 108 | making modifications to it. "Object code" means any non-source form of 109 | a work. 110 | 111 | A "Standard Interface" means an interface that either is an official 112 | standard defined by a recognized standards body, or, in the case of 113 | interfaces specified for a particular programming language, one that 114 | is widely used among developers working in that language. 115 | 116 | The "System Libraries" of an executable work include anything, other 117 | than the work as a whole, that (a) is included in the normal form of 118 | packaging a Major Component, but which is not part of that Major 119 | Component, and (b) serves only to enable use of the work with that 120 | Major Component, or to implement a Standard Interface for which an 121 | implementation is available to the public in source code form. A 122 | "Major Component", in this context, means a major essential component 123 | (kernel, window system, and so on) of the specific operating system 124 | (if any) on which the executable work runs, or a compiler used to 125 | produce the work, or an object code interpreter used to run it. 126 | 127 | The "Corresponding Source" for a work in object code form means all 128 | the source code needed to generate, install, and (for an executable 129 | work) run the object code and to modify the work, including scripts to 130 | control those activities. However, it does not include the work's 131 | System Libraries, or general-purpose tools or generally available free 132 | programs which are used unmodified in performing those activities but 133 | which are not part of the work. For example, Corresponding Source 134 | includes interface definition files associated with source files for 135 | the work, and the source code for shared libraries and dynamically 136 | linked subprograms that the work is specifically designed to require, 137 | such as by intimate data communication or control flow between those 138 | subprograms and other parts of the work. 139 | 140 | The Corresponding Source need not include anything that users can 141 | regenerate automatically from other parts of the Corresponding Source. 142 | 143 | The Corresponding Source for a work in source code form is that same 144 | work. 145 | 146 | ### 2. Basic Permissions. 147 | 148 | All rights granted under this License are granted for the term of 149 | copyright on the Program, and are irrevocable provided the stated 150 | conditions are met. This License explicitly affirms your unlimited 151 | permission to run the unmodified Program. The output from running a 152 | covered work is covered by this License only if the output, given its 153 | content, constitutes a covered work. This License acknowledges your 154 | rights of fair use or other equivalent, as provided by copyright law. 155 | 156 | You may make, run and propagate covered works that you do not convey, 157 | without conditions so long as your license otherwise remains in force. 158 | You may convey covered works to others for the sole purpose of having 159 | them make modifications exclusively for you, or provide you with 160 | facilities for running those works, provided that you comply with the 161 | terms of this License in conveying all material for which you do not 162 | control copyright. Those thus making or running the covered works for 163 | you must do so exclusively on your behalf, under your direction and 164 | control, on terms that prohibit them from making any copies of your 165 | copyrighted material outside their relationship with you. 166 | 167 | Conveying under any other circumstances is permitted solely under the 168 | conditions stated below. Sublicensing is not allowed; section 10 makes 169 | it unnecessary. 170 | 171 | ### 3. Protecting Users' Legal Rights From Anti-Circumvention Law. 172 | 173 | No covered work shall be deemed part of an effective technological 174 | measure under any applicable law fulfilling obligations under article 175 | 11 of the WIPO copyright treaty adopted on 20 December 1996, or 176 | similar laws prohibiting or restricting circumvention of such 177 | measures. 178 | 179 | When you convey a covered work, you waive any legal power to forbid 180 | circumvention of technological measures to the extent such 181 | circumvention is effected by exercising rights under this License with 182 | respect to the covered work, and you disclaim any intention to limit 183 | operation or modification of the work as a means of enforcing, against 184 | the work's users, your or third parties' legal rights to forbid 185 | circumvention of technological measures. 186 | 187 | ### 4. Conveying Verbatim Copies. 188 | 189 | You may convey verbatim copies of the Program's source code as you 190 | receive it, in any medium, provided that you conspicuously and 191 | appropriately publish on each copy an appropriate copyright notice; 192 | keep intact all notices stating that this License and any 193 | non-permissive terms added in accord with section 7 apply to the code; 194 | keep intact all notices of the absence of any warranty; and give all 195 | recipients a copy of this License along with the Program. 196 | 197 | You may charge any price or no price for each copy that you convey, 198 | and you may offer support or warranty protection for a fee. 199 | 200 | ### 5. Conveying Modified Source Versions. 201 | 202 | You may convey a work based on the Program, or the modifications to 203 | produce it from the Program, in the form of source code under the 204 | terms of section 4, provided that you also meet all of these 205 | conditions: 206 | 207 | - a) The work must carry prominent notices stating that you modified 208 | it, and giving a relevant date. 209 | - b) The work must carry prominent notices stating that it is 210 | released under this License and any conditions added under 211 | section 7. This requirement modifies the requirement in section 4 212 | to "keep intact all notices". 213 | - c) You must license the entire work, as a whole, under this 214 | License to anyone who comes into possession of a copy. This 215 | License will therefore apply, along with any applicable section 7 216 | additional terms, to the whole of the work, and all its parts, 217 | regardless of how they are packaged. This License gives no 218 | permission to license the work in any other way, but it does not 219 | invalidate such permission if you have separately received it. 220 | - d) If the work has interactive user interfaces, each must display 221 | Appropriate Legal Notices; however, if the Program has interactive 222 | interfaces that do not display Appropriate Legal Notices, your 223 | work need not make them do so. 224 | 225 | A compilation of a covered work with other separate and independent 226 | works, which are not by their nature extensions of the covered work, 227 | and which are not combined with it such as to form a larger program, 228 | in or on a volume of a storage or distribution medium, is called an 229 | "aggregate" if the compilation and its resulting copyright are not 230 | used to limit the access or legal rights of the compilation's users 231 | beyond what the individual works permit. Inclusion of a covered work 232 | in an aggregate does not cause this License to apply to the other 233 | parts of the aggregate. 234 | 235 | ### 6. Conveying Non-Source Forms. 236 | 237 | You may convey a covered work in object code form under the terms of 238 | sections 4 and 5, provided that you also convey the machine-readable 239 | Corresponding Source under the terms of this License, in one of these 240 | ways: 241 | 242 | - a) Convey the object code in, or embodied in, a physical product 243 | (including a physical distribution medium), accompanied by the 244 | Corresponding Source fixed on a durable physical medium 245 | customarily used for software interchange. 246 | - b) Convey the object code in, or embodied in, a physical product 247 | (including a physical distribution medium), accompanied by a 248 | written offer, valid for at least three years and valid for as 249 | long as you offer spare parts or customer support for that product 250 | model, to give anyone who possesses the object code either (1) a 251 | copy of the Corresponding Source for all the software in the 252 | product that is covered by this License, on a durable physical 253 | medium customarily used for software interchange, for a price no 254 | more than your reasonable cost of physically performing this 255 | conveying of source, or (2) access to copy the Corresponding 256 | Source from a network server at no charge. 257 | - c) Convey individual copies of the object code with a copy of the 258 | written offer to provide the Corresponding Source. This 259 | alternative is allowed only occasionally and noncommercially, and 260 | only if you received the object code with such an offer, in accord 261 | with subsection 6b. 262 | - d) Convey the object code by offering access from a designated 263 | place (gratis or for a charge), and offer equivalent access to the 264 | Corresponding Source in the same way through the same place at no 265 | further charge. You need not require recipients to copy the 266 | Corresponding Source along with the object code. If the place to 267 | copy the object code is a network server, the Corresponding Source 268 | may be on a different server (operated by you or a third party) 269 | that supports equivalent copying facilities, provided you maintain 270 | clear directions next to the object code saying where to find the 271 | Corresponding Source. Regardless of what server hosts the 272 | Corresponding Source, you remain obligated to ensure that it is 273 | available for as long as needed to satisfy these requirements. 274 | - e) Convey the object code using peer-to-peer transmission, 275 | provided you inform other peers where the object code and 276 | Corresponding Source of the work are being offered to the general 277 | public at no charge under subsection 6d. 278 | 279 | A separable portion of the object code, whose source code is excluded 280 | from the Corresponding Source as a System Library, need not be 281 | included in conveying the object code work. 282 | 283 | A "User Product" is either (1) a "consumer product", which means any 284 | tangible personal property which is normally used for personal, 285 | family, or household purposes, or (2) anything designed or sold for 286 | incorporation into a dwelling. In determining whether a product is a 287 | consumer product, doubtful cases shall be resolved in favor of 288 | coverage. For a particular product received by a particular user, 289 | "normally used" refers to a typical or common use of that class of 290 | product, regardless of the status of the particular user or of the way 291 | in which the particular user actually uses, or expects or is expected 292 | to use, the product. A product is a consumer product regardless of 293 | whether the product has substantial commercial, industrial or 294 | non-consumer uses, unless such uses represent the only significant 295 | mode of use of the product. 296 | 297 | "Installation Information" for a User Product means any methods, 298 | procedures, authorization keys, or other information required to 299 | install and execute modified versions of a covered work in that User 300 | Product from a modified version of its Corresponding Source. The 301 | information must suffice to ensure that the continued functioning of 302 | the modified object code is in no case prevented or interfered with 303 | solely because modification has been made. 304 | 305 | If you convey an object code work under this section in, or with, or 306 | specifically for use in, a User Product, and the conveying occurs as 307 | part of a transaction in which the right of possession and use of the 308 | User Product is transferred to the recipient in perpetuity or for a 309 | fixed term (regardless of how the transaction is characterized), the 310 | Corresponding Source conveyed under this section must be accompanied 311 | by the Installation Information. But this requirement does not apply 312 | if neither you nor any third party retains the ability to install 313 | modified object code on the User Product (for example, the work has 314 | been installed in ROM). 315 | 316 | The requirement to provide Installation Information does not include a 317 | requirement to continue to provide support service, warranty, or 318 | updates for a work that has been modified or installed by the 319 | recipient, or for the User Product in which it has been modified or 320 | installed. Access to a network may be denied when the modification 321 | itself materially and adversely affects the operation of the network 322 | or violates the rules and protocols for communication across the 323 | network. 324 | 325 | Corresponding Source conveyed, and Installation Information provided, 326 | in accord with this section must be in a format that is publicly 327 | documented (and with an implementation available to the public in 328 | source code form), and must require no special password or key for 329 | unpacking, reading or copying. 330 | 331 | ### 7. Additional Terms. 332 | 333 | "Additional permissions" are terms that supplement the terms of this 334 | License by making exceptions from one or more of its conditions. 335 | Additional permissions that are applicable to the entire Program shall 336 | be treated as though they were included in this License, to the extent 337 | that they are valid under applicable law. If additional permissions 338 | apply only to part of the Program, that part may be used separately 339 | under those permissions, but the entire Program remains governed by 340 | this License without regard to the additional permissions. 341 | 342 | When you convey a copy of a covered work, you may at your option 343 | remove any additional permissions from that copy, or from any part of 344 | it. (Additional permissions may be written to require their own 345 | removal in certain cases when you modify the work.) You may place 346 | additional permissions on material, added by you to a covered work, 347 | for which you have or can give appropriate copyright permission. 348 | 349 | Notwithstanding any other provision of this License, for material you 350 | add to a covered work, you may (if authorized by the copyright holders 351 | of that material) supplement the terms of this License with terms: 352 | 353 | - a) Disclaiming warranty or limiting liability differently from the 354 | terms of sections 15 and 16 of this License; or 355 | - b) Requiring preservation of specified reasonable legal notices or 356 | author attributions in that material or in the Appropriate Legal 357 | Notices displayed by works containing it; or 358 | - c) Prohibiting misrepresentation of the origin of that material, 359 | or requiring that modified versions of such material be marked in 360 | reasonable ways as different from the original version; or 361 | - d) Limiting the use for publicity purposes of names of licensors 362 | or authors of the material; or 363 | - e) Declining to grant rights under trademark law for use of some 364 | trade names, trademarks, or service marks; or 365 | - f) Requiring indemnification of licensors and authors of that 366 | material by anyone who conveys the material (or modified versions 367 | of it) with contractual assumptions of liability to the recipient, 368 | for any liability that these contractual assumptions directly 369 | impose on those licensors and authors. 370 | 371 | All other non-permissive additional terms are considered "further 372 | restrictions" within the meaning of section 10. If the Program as you 373 | received it, or any part of it, contains a notice stating that it is 374 | governed by this License along with a term that is a further 375 | restriction, you may remove that term. If a license document contains 376 | a further restriction but permits relicensing or conveying under this 377 | License, you may add to a covered work material governed by the terms 378 | of that license document, provided that the further restriction does 379 | not survive such relicensing or conveying. 380 | 381 | If you add terms to a covered work in accord with this section, you 382 | must place, in the relevant source files, a statement of the 383 | additional terms that apply to those files, or a notice indicating 384 | where to find the applicable terms. 385 | 386 | Additional terms, permissive or non-permissive, may be stated in the 387 | form of a separately written license, or stated as exceptions; the 388 | above requirements apply either way. 389 | 390 | ### 8. Termination. 391 | 392 | You may not propagate or modify a covered work except as expressly 393 | provided under this License. Any attempt otherwise to propagate or 394 | modify it is void, and will automatically terminate your rights under 395 | this License (including any patent licenses granted under the third 396 | paragraph of section 11). 397 | 398 | However, if you cease all violation of this License, then your license 399 | from a particular copyright holder is reinstated (a) provisionally, 400 | unless and until the copyright holder explicitly and finally 401 | terminates your license, and (b) permanently, if the copyright holder 402 | fails to notify you of the violation by some reasonable means prior to 403 | 60 days after the cessation. 404 | 405 | Moreover, your license from a particular copyright holder is 406 | reinstated permanently if the copyright holder notifies you of the 407 | violation by some reasonable means, this is the first time you have 408 | received notice of violation of this License (for any work) from that 409 | copyright holder, and you cure the violation prior to 30 days after 410 | your receipt of the notice. 411 | 412 | Termination of your rights under this section does not terminate the 413 | licenses of parties who have received copies or rights from you under 414 | this License. If your rights have been terminated and not permanently 415 | reinstated, you do not qualify to receive new licenses for the same 416 | material under section 10. 417 | 418 | ### 9. Acceptance Not Required for Having Copies. 419 | 420 | You are not required to accept this License in order to receive or run 421 | a copy of the Program. Ancillary propagation of a covered work 422 | occurring solely as a consequence of using peer-to-peer transmission 423 | to receive a copy likewise does not require acceptance. However, 424 | nothing other than this License grants you permission to propagate or 425 | modify any covered work. These actions infringe copyright if you do 426 | not accept this License. Therefore, by modifying or propagating a 427 | covered work, you indicate your acceptance of this License to do so. 428 | 429 | ### 10. Automatic Licensing of Downstream Recipients. 430 | 431 | Each time you convey a covered work, the recipient automatically 432 | receives a license from the original licensors, to run, modify and 433 | propagate that work, subject to this License. You are not responsible 434 | for enforcing compliance by third parties with this License. 435 | 436 | An "entity transaction" is a transaction transferring control of an 437 | organization, or substantially all assets of one, or subdividing an 438 | organization, or merging organizations. If propagation of a covered 439 | work results from an entity transaction, each party to that 440 | transaction who receives a copy of the work also receives whatever 441 | licenses to the work the party's predecessor in interest had or could 442 | give under the previous paragraph, plus a right to possession of the 443 | Corresponding Source of the work from the predecessor in interest, if 444 | the predecessor has it or can get it with reasonable efforts. 445 | 446 | You may not impose any further restrictions on the exercise of the 447 | rights granted or affirmed under this License. For example, you may 448 | not impose a license fee, royalty, or other charge for exercise of 449 | rights granted under this License, and you may not initiate litigation 450 | (including a cross-claim or counterclaim in a lawsuit) alleging that 451 | any patent claim is infringed by making, using, selling, offering for 452 | sale, or importing the Program or any portion of it. 453 | 454 | ### 11. Patents. 455 | 456 | A "contributor" is a copyright holder who authorizes use under this 457 | License of the Program or a work on which the Program is based. The 458 | work thus licensed is called the contributor's "contributor version". 459 | 460 | A contributor's "essential patent claims" are all patent claims owned 461 | or controlled by the contributor, whether already acquired or 462 | hereafter acquired, that would be infringed by some manner, permitted 463 | by this License, of making, using, or selling its contributor version, 464 | but do not include claims that would be infringed only as a 465 | consequence of further modification of the contributor version. For 466 | purposes of this definition, "control" includes the right to grant 467 | patent sublicenses in a manner consistent with the requirements of 468 | this License. 469 | 470 | Each contributor grants you a non-exclusive, worldwide, royalty-free 471 | patent license under the contributor's essential patent claims, to 472 | make, use, sell, offer for sale, import and otherwise run, modify and 473 | propagate the contents of its contributor version. 474 | 475 | In the following three paragraphs, a "patent license" is any express 476 | agreement or commitment, however denominated, not to enforce a patent 477 | (such as an express permission to practice a patent or covenant not to 478 | sue for patent infringement). To "grant" such a patent license to a 479 | party means to make such an agreement or commitment not to enforce a 480 | patent against the party. 481 | 482 | If you convey a covered work, knowingly relying on a patent license, 483 | and the Corresponding Source of the work is not available for anyone 484 | to copy, free of charge and under the terms of this License, through a 485 | publicly available network server or other readily accessible means, 486 | then you must either (1) cause the Corresponding Source to be so 487 | available, or (2) arrange to deprive yourself of the benefit of the 488 | patent license for this particular work, or (3) arrange, in a manner 489 | consistent with the requirements of this License, to extend the patent 490 | license to downstream recipients. "Knowingly relying" means you have 491 | actual knowledge that, but for the patent license, your conveying the 492 | covered work in a country, or your recipient's use of the covered work 493 | in a country, would infringe one or more identifiable patents in that 494 | country that you have reason to believe are valid. 495 | 496 | If, pursuant to or in connection with a single transaction or 497 | arrangement, you convey, or propagate by procuring conveyance of, a 498 | covered work, and grant a patent license to some of the parties 499 | receiving the covered work authorizing them to use, propagate, modify 500 | or convey a specific copy of the covered work, then the patent license 501 | you grant is automatically extended to all recipients of the covered 502 | work and works based on it. 503 | 504 | A patent license is "discriminatory" if it does not include within the 505 | scope of its coverage, prohibits the exercise of, or is conditioned on 506 | the non-exercise of one or more of the rights that are specifically 507 | granted under this License. You may not convey a covered work if you 508 | are a party to an arrangement with a third party that is in the 509 | business of distributing software, under which you make payment to the 510 | third party based on the extent of your activity of conveying the 511 | work, and under which the third party grants, to any of the parties 512 | who would receive the covered work from you, a discriminatory patent 513 | license (a) in connection with copies of the covered work conveyed by 514 | you (or copies made from those copies), or (b) primarily for and in 515 | connection with specific products or compilations that contain the 516 | covered work, unless you entered into that arrangement, or that patent 517 | license was granted, prior to 28 March 2007. 518 | 519 | Nothing in this License shall be construed as excluding or limiting 520 | any implied license or other defenses to infringement that may 521 | otherwise be available to you under applicable patent law. 522 | 523 | ### 12. No Surrender of Others' Freedom. 524 | 525 | If conditions are imposed on you (whether by court order, agreement or 526 | otherwise) that contradict the conditions of this License, they do not 527 | excuse you from the conditions of this License. If you cannot convey a 528 | covered work so as to satisfy simultaneously your obligations under 529 | this License and any other pertinent obligations, then as a 530 | consequence you may not convey it at all. For example, if you agree to 531 | terms that obligate you to collect a royalty for further conveying 532 | from those to whom you convey the Program, the only way you could 533 | satisfy both those terms and this License would be to refrain entirely 534 | from conveying the Program. 535 | 536 | ### 13. Remote Network Interaction; Use with the GNU General Public License. 537 | 538 | Notwithstanding any other provision of this License, if you modify the 539 | Program, your modified version must prominently offer all users 540 | interacting with it remotely through a computer network (if your 541 | version supports such interaction) an opportunity to receive the 542 | Corresponding Source of your version by providing access to the 543 | Corresponding Source from a network server at no charge, through some 544 | standard or customary means of facilitating copying of software. This 545 | Corresponding Source shall include the Corresponding Source for any 546 | work covered by version 3 of the GNU General Public License that is 547 | incorporated pursuant to the following paragraph. 548 | 549 | Notwithstanding any other provision of this License, you have 550 | permission to link or combine any covered work with a work licensed 551 | under version 3 of the GNU General Public License into a single 552 | combined work, and to convey the resulting work. The terms of this 553 | License will continue to apply to the part which is the covered work, 554 | but the work with which it is combined will remain governed by version 555 | 3 of the GNU General Public License. 556 | 557 | ### 14. Revised Versions of this License. 558 | 559 | The Free Software Foundation may publish revised and/or new versions 560 | of the GNU Affero General Public License from time to time. Such new 561 | versions will be similar in spirit to the present version, but may 562 | differ in detail to address new problems or concerns. 563 | 564 | Each version is given a distinguishing version number. If the Program 565 | specifies that a certain numbered version of the GNU Affero General 566 | Public License "or any later version" applies to it, you have the 567 | option of following the terms and conditions either of that numbered 568 | version or of any later version published by the Free Software 569 | Foundation. If the Program does not specify a version number of the 570 | GNU Affero General Public License, you may choose any version ever 571 | published by the Free Software Foundation. 572 | 573 | If the Program specifies that a proxy can decide which future versions 574 | of the GNU Affero General Public License can be used, that proxy's 575 | public statement of acceptance of a version permanently authorizes you 576 | to choose that version for the Program. 577 | 578 | Later license versions may give you additional or different 579 | permissions. However, no additional obligations are imposed on any 580 | author or copyright holder as a result of your choosing to follow a 581 | later version. 582 | 583 | ### 15. Disclaimer of Warranty. 584 | 585 | THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY 586 | APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT 587 | HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT 588 | WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT 589 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 590 | A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND 591 | PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE 592 | DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, REPAIR OR 593 | CORRECTION. 594 | 595 | ### 16. Limitation of Liability. 596 | 597 | IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING 598 | WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR 599 | CONVEYS THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, 600 | INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES 601 | ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT 602 | NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR 603 | LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM 604 | TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER 605 | PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. 606 | 607 | ### 17. Interpretation of Sections 15 and 16. 608 | 609 | If the disclaimer of warranty and limitation of liability provided 610 | above cannot be given local legal effect according to their terms, 611 | reviewing courts shall apply local law that most closely approximates 612 | an absolute waiver of all civil liability in connection with the 613 | Program, unless a warranty or assumption of liability accompanies a 614 | copy of the Program in return for a fee. 615 | 616 | END OF TERMS AND CONDITIONS 617 | 618 | ## How to Apply These Terms to Your New Programs 619 | 620 | If you develop a new program, and you want it to be of the greatest 621 | possible use to the public, the best way to achieve this is to make it 622 | free software which everyone can redistribute and change under these 623 | terms. 624 | 625 | To do so, attach the following notices to the program. It is safest to 626 | attach them to the start of each source file to most effectively state 627 | the exclusion of warranty; and each file should have at least the 628 | "copyright" line and a pointer to where the full notice is found. 629 | 630 | 631 | Copyright (C) 632 | 633 | This program is free software: you can redistribute it and/or modify 634 | it under the terms of the GNU Affero General Public License as 635 | published by the Free Software Foundation, either version 3 of the 636 | License, or (at your option) any later version. 637 | 638 | This program is distributed in the hope that it will be useful, 639 | but WITHOUT ANY WARRANTY; without even the implied warranty of 640 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 641 | GNU Affero General Public License for more details. 642 | 643 | You should have received a copy of the GNU Affero General Public License 644 | along with this program. If not, see . 645 | 646 | Also add information on how to contact you by electronic and paper 647 | mail. 648 | 649 | If your software can interact with users remotely through a computer 650 | network, you should also make sure that it provides a way for users to 651 | get its source. For example, if your program is a web application, its 652 | interface could display a "Source" link that leads users to an archive 653 | of the code. There are many ways you could offer source, and different 654 | solutions will be better for different programs; see section 13 for 655 | the specific requirements. 656 | 657 | You should also get your employer (if you work as a programmer) or 658 | school, if any, to sign a "copyright disclaimer" for the program, if 659 | necessary. For more information on this, and how to apply and follow 660 | the GNU AGPL, see . 661 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Zebra 2 | A vector database for querying meaningfully similar data. 3 | 4 | ## Features 5 | * **On-disk database index**, minimising memory impact for large datasets 6 | * Distance metrics and embedding models are provided, though external implementations can be supplied 7 | * Inserting, **deleting**, and querying vectors 8 | * **[No memory-mapped (MMAP) file IO](https://db.cs.cmu.edu/mmap-cidr2022/)** 9 | * Parallelised database operations; multithreaded reads & writes are safe 10 | 11 | ### Motivation 12 | Approximate nearest neighbour search for finding semantically-similar documents is a common use case, and a variety of existing solutions exist that are often described as 'vector databases'. Many of these solutions are offered as services, and some exist as libraries. 13 | 14 | While designing a content recommendation system for [Oku](https://okubrowser.github.io), several requirements became clear: 15 | * Safe multithreaded database access 16 | * Inserting, deleting, and querying records 17 | * On-disk storage of database contents 18 | * Support for multiple [modalities](https://en.wikipedia.org/wiki/Modality_(semiotics)) of information 19 | 20 | Oku enables distributed storage & distribution of mutable user-generated data; therefore, any embedded database used for semantic search needed to be: 21 | * Capable of asynchronous read & write access 22 | * Capable of creating, reading, updating, and deleting (**[CRUD](https://en.wikipedia.org/wiki/Create,_read,_update_and_delete)**) records **without downtime or excessive resource usage** 23 | * Capable of scaling in dataset size **without excessively impacting memory usage for an individual node** 24 | * Capable of acceptable recall with multiple distance metrics 25 | 26 | Despite the common need—a scalable CRUD database—existing solutions often fell short. 27 | 28 | #### Distribution & CRUD 29 | Many vector databases utilise the [hierarchical navigable small world (HNSW)](https://en.wikipedia.org/wiki/Hierarchical_navigable_small_world) algorithm to construct their database indices, as it (a) achieves high recall on high-dimensional data regardless of distance metric, and (b) performs fast queries regardless of dataset size. However, despite its attractiveness on benchmarks, it can be impractical to use in many production contexts as (a) it's difficult to distribute as you cannot shard the index across multiple nodes, (b) the entire index must be loaded into memory to perform operations, making memory a bottleneck in addition to storage, and (c) deleting vectors essentially requires rebuilding the entire index from scratch and re-inserting every vector except the deleted ones; using redundant indices and tombstoning is the only way to keep the database online. 30 | 31 | The need for a scalable & mutable vector database is not new, however, and the problem has apparently been solved to an acceptable degree before—content recommendation systems based on embedding vectors have been in production for many years, and they've often used some variation of [locality sensitive hashing (LSH)](https://en.wikipedia.org/wiki/Locality-sensitive_hashing) to build a vector database index. An LSH index is not graph-based, but instead breaks up an *f*-dimensional space into regions of similar vectors. Consequently, it can be (a) sharded, (b) accessed in parallel, and (c) accessed from storage because, unlike a graph such as HNSW, it is not '[object soup](https://jacko.io/object_soup.html)' and avoids issues with cache locality and synchronisation in multithreaded contexts. LSH's advantages in performance and resource usage does come with an implication: while HNSW approximates neighbours, LSH approximates similarities. The recall of LSH is lesser as it's less concerned with finding *the nearest* neighbours, and more concerned with just finding what *is near*. For fine-grained searches, LSH is less helpful, but for large & varied datasets where it's important to find records that are 'close enough', it has significant advantages. 32 | 33 | #### Integrity & Safety 34 | To avoid excessive memory usage, some have saved indexes to storage and performed operations directly on the index files as if they were in memory, taking advantage of a technique called [memory mapping (`mmap`)](https://en.wikipedia.org/wiki/Memory-mapped_file). Spotify boasts of [its LSH index](https://github.com/spotify/annoy): 35 | > … you can share index across processes … you can pass around indexes as files and map them into memory quickly … You can also pass around and distribute static files to use in production environment, in Hadoop jobs, etc. Any process will be able to load (mmap) the index into memory and will be able to do lookups immediately. 36 | 37 | Cloudflare [makes similar bold claims](https://blog.cloudflare.com/scalable-machine-learning-at-cloudflare/) regarding its use of `mmap`: 38 | > Leveraging the benefits of memory-mapped files, wait-free synchronization and zero-copy deserialization, we've crafted a unique and powerful tool for managing high-performance, concurrent data access between processes. 39 | > 40 | > The data is stored in shared mapped memory, which allows the `Synchronizer` to “write” to it and “read” from it concurrently. This design makes `mmap-sync` a highly efficient and flexible tool for managing shared, concurrent data access. 41 | > 42 | > In the wake of our redesign, we've constructed a powerful and efficient system that truly embodies the essence of 'bliss'. Harnessing the advantages of memory-mapped files, wait-free synchronization, allocation-free operations, and zero-copy deserialization, we've established a robust infrastructure that maintains peak performance while achieving remarkable reductions in latency. 43 | 44 | Multiprocess concurrency with `mmap` is arguably *impossible*—countless DBMSes have learned the same lesson after many years. There is [a paper on this subject](https://db.cs.cmu.edu/papers/2022/cidr2022-p13-crotty.pdf) that covers the pitfalls of `mmap` in detail; suffice it to say, a memory-mapped database index has not demonstrably achieved the data integrity & memory-safety guarantees necessary for a production database. 45 | 46 | #### Potential Improvements 47 | This software is free & open-source (FOSS), and code contributions are welcome. 48 | Its usage within Oku involves operating on consumer hardware with data that is diverse. For use-cases where more hardware resources are available or datasets needing greater recall are used, this database could be extended with a new implementation of the HNSW algorithm, providing a mutable, multithreaded, real-time variant of HNSW. Such an index may be possible, as [Vespa claims to have created such an implementation](https://docs.vespa.ai/en/approximate-nn-hnsw.html). Improving HNSW's memory usage, however, appears impossible due to its fundamental graph-based nature. 49 | 50 | ##### The name 51 | Zebras make a 'neigh' sound. The database performs an [approximate nearest-*neigh*bour search](https://en.wikipedia.org/wiki/Nearest_neighbor_search#Approximation_methods) to find similar data. 52 | 53 | ## Installation 54 | Zebra is intended for use as an embedded database. You can add it as a dependency to a Rust project with the following command: 55 | ```sh 56 | cargo add --git "https://github.com/emmyoh/zebra" 57 | ``` 58 | 59 | Additionally, a command-line interface (CLI) exists for basic usage in the terminal. 60 | 61 | With the [Rust toolchain](https://rustup.rs/) installed, run the following command: 62 | ```sh 63 | cargo install --git https://github.com/emmyoh/zebra --features="cli" 64 | ``` 65 | 66 | You should specify the features relevant to your use case. For example, if you're interested in using the Zebra CLI on an Apple silicon device, run: 67 | ```sh 68 | cargo install --git https://github.com/emmyoh/zebra --features="cli,accelerate,metal" 69 | ``` 70 | 71 | ### Features 72 | * `default_db` - Provides default configurations for databases. 73 | * `accelerate` - Uses Apple's Accelerate framework when running on Apple operating systems. 74 | * `cuda` - Enables GPU support with Nvidia cards. 75 | * `mkl` - Uses Intel oneMKL with Intel CPUs and GPUs. 76 | * `metal` - Enables GPU support for Apple silicon machines. 77 | * `sixel` - Prints images in Sixel format when using the CLI with compatible terminals. 78 | * `cli` - Provides a command-line interface to Zebra. -------------------------------------------------------------------------------- /_rust-toolchain.toml: -------------------------------------------------------------------------------- 1 | [toolchain] 2 | channel = "nightly" 3 | -------------------------------------------------------------------------------- /precommit.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | cargo clippy --fix --allow-dirty --features="cli" 3 | # __CARGO_FIX_YOLO=1 cargo clippy --fix --broken-code --allow-dirty --features="cli" 4 | cargo fmt 5 | cargo check 6 | cargo check --features="cli" 7 | -------------------------------------------------------------------------------- /src/database/core.rs: -------------------------------------------------------------------------------- 1 | use super::index::lsh::{LSHIndex, LSHIndexOptions}; 2 | use crate::Embedding; 3 | use crate::{distance::DistanceUnit, model::core::DatabaseEmbeddingModel}; 4 | use bytes::Bytes; 5 | use dashmap::{DashMap, DashSet}; 6 | use rayon::iter::IntoParallelIterator; 7 | use rayon::iter::ParallelIterator; 8 | use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator}; 9 | use rayon::slice::ParallelSliceMut; 10 | use serde::{Deserialize, Serialize}; 11 | use space::Metric; 12 | use std::io::Cursor; 13 | use std::{ 14 | fs::{self, OpenOptions}, 15 | io::{self, BufReader, BufWriter}, 16 | }; 17 | use uuid::Uuid; 18 | 19 | #[derive(Debug, Clone, Serialize, Deserialize)] 20 | struct DatabaseInner< 21 | const N: usize, 22 | Met: Metric, Unit = DistanceUnit> + Default + Serialize + Send + Sync, 23 | Mod: DatabaseEmbeddingModel + Default + Serialize + Send + Sync, 24 | > { 25 | uuid: Uuid, 26 | model: Mod, 27 | metric: Met, 28 | index_options: LSHIndexOptions, 29 | } 30 | 31 | impl< 32 | const N: usize, 33 | Met: Metric, Unit = DistanceUnit> + Default + Serialize + Send + Sync, 34 | Mod: DatabaseEmbeddingModel + Default + Serialize + Send + Sync, 35 | > DatabaseInner 36 | where 37 | for<'de> Mod: Deserialize<'de>, 38 | for<'de> Met: Deserialize<'de>, 39 | { 40 | fn index(&self) -> anyhow::Result> { 41 | LSHIndex::new(&self.uuid, &self.index_options) 42 | } 43 | } 44 | 45 | #[derive(Clone)] 46 | /// A database containing embedding vectors and documents. 47 | /// 48 | /// # Arguments 49 | /// 50 | /// * `N` - The dimensionality of the vectors in the database. 51 | /// 52 | /// * `Met` - The distance metric used by the database index. 53 | /// 54 | /// * `Mod` - The model used to generate the embedding vectors. 55 | pub struct Database< 56 | const N: usize, 57 | Met: Metric, Unit = DistanceUnit> + Default + Serialize + Send + Sync, 58 | Mod: DatabaseEmbeddingModel + Default + Serialize + Send + Sync, 59 | > { 60 | inner: DatabaseInner, 61 | /// The database index used to approximate nearest-neighbour search. 62 | pub index: LSHIndex, 63 | path: String, 64 | } 65 | 66 | impl< 67 | const N: usize, 68 | Met: Metric, Unit = DistanceUnit> + Default + Serialize + Send + Sync, 69 | Mod: DatabaseEmbeddingModel + Default + Serialize + Send + Sync, 70 | > Database 71 | where 72 | for<'de> Mod: Deserialize<'de>, 73 | for<'de> Met: Deserialize<'de>, 74 | { 75 | fn database_subdirectory(&self) -> String { 76 | format!("{}", self.inner.uuid.as_simple()) 77 | } 78 | 79 | fn default_database_path(&self) -> String { 80 | format!("{}.zebra", self.inner.uuid.as_simple()) 81 | } 82 | 83 | /// Load the database from disk. 84 | /// 85 | /// # Arguments 86 | /// 87 | /// * `path` - The path to the database file. 88 | /// 89 | /// # Returns 90 | /// 91 | /// A [Database] containing embeddings & documents. 92 | pub fn open(path: &String) -> anyhow::Result { 93 | let db_bytes = fs::read(path)?; 94 | let inner: DatabaseInner = 95 | bincode::serde::decode_from_slice(&db_bytes, bincode::config::legacy())?.0; 96 | let index = inner.index()?; 97 | Ok(Self { 98 | inner, 99 | path: path.clone(), 100 | index, 101 | }) 102 | } 103 | 104 | /// Create a database in memory.\ 105 | /// Note: All database operations require read & write access to storage; a database cannot be used only in memory. 106 | /// 107 | /// # Returns 108 | /// 109 | /// A new [Database]. 110 | pub fn new(index_options: &LSHIndexOptions) -> anyhow::Result { 111 | let uuid = Uuid::now_v7(); 112 | let inner = DatabaseInner { 113 | uuid, 114 | model: Mod::default(), 115 | metric: Met::default(), 116 | index_options: index_options.clone(), 117 | }; 118 | let index = inner.index()?; 119 | let mut new = Self { 120 | inner, 121 | index, 122 | path: String::new(), 123 | }; 124 | new.path = new.default_database_path(); 125 | new.save_database(None)?; 126 | Ok(new) 127 | } 128 | 129 | /// Create a database in memory, persisting to storage at the given path. 130 | /// 131 | /// # Arguments 132 | /// 133 | /// * `path` - The path to the database file. 134 | /// 135 | /// # Returns 136 | /// 137 | /// A [Database] containing embeddings & documents. 138 | pub fn new_with_path( 139 | path: &String, 140 | index_options: &LSHIndexOptions, 141 | ) -> anyhow::Result { 142 | let uuid = Uuid::now_v7(); 143 | let inner = DatabaseInner { 144 | uuid, 145 | model: Mod::default(), 146 | metric: Met::default(), 147 | index_options: index_options.clone(), 148 | }; 149 | let index = inner.index()?; 150 | let new = Self { 151 | inner, 152 | index, 153 | path: path.to_owned(), 154 | }; 155 | new.save_database(None)?; 156 | Ok(new) 157 | } 158 | 159 | /// Load the database from disk, or create it if it does not already exist. 160 | /// 161 | /// # Arguments 162 | /// 163 | /// * `path` - The path to the database file. 164 | /// 165 | /// # Returns 166 | /// 167 | /// A [Database] containing embeddings & documents. 168 | pub fn open_or_create( 169 | path: &String, 170 | index_options: &LSHIndexOptions, 171 | ) -> anyhow::Result { 172 | match Self::open(path) { 173 | Ok(db) => Ok(db), 174 | Err(_) => Ok(Self::new_with_path(path, index_options)?), 175 | } 176 | } 177 | 178 | /// Save the database to disk. 179 | /// 180 | /// # Arguments 181 | /// 182 | /// * `path` - An optional path to save the database to; if left blank, will use the path the database was opened from. 183 | pub fn save_database(&self, path: Option<&String>) -> anyhow::Result<()> { 184 | fs::write( 185 | path.unwrap_or(&self.path), 186 | bincode::serde::encode_to_vec(&self.inner, bincode::config::legacy())?, 187 | )?; 188 | self.index.save()?; 189 | Ok(()) 190 | } 191 | 192 | /// Delete the database and its contents, including all vectors and documents.\ 193 | /// Note: This deletes the file at the path the database was opened from; if the database file was moved after opening, this may have unintended consequences. 194 | pub fn clear_database(&self) { 195 | let _ = self.index.clear(); 196 | let _ = std::fs::remove_file(&self.path); 197 | let _ = std::fs::remove_dir_all(self.database_subdirectory()); 198 | } 199 | 200 | /// Removes records from the database. 201 | /// 202 | /// # Arguments 203 | /// 204 | /// * `embedding_ids` - The IDs of the vectors to remove. 205 | pub fn remove(&self, embedding_ids: &Vec) -> anyhow::Result<()> { 206 | let document_subdirectory = self.database_subdirectory(); 207 | let removed = self.index.remove(embedding_ids)?; 208 | removed.into_par_iter().for_each(|x| { 209 | let _ = std::fs::remove_file(format!("{}/{}.lz4", document_subdirectory, x)); 210 | }); 211 | self.save_database(None)?; 212 | Ok(()) 213 | } 214 | 215 | /// Remove duplicate embedding vectors from the database. 216 | pub fn deduplicate(&self) -> anyhow::Result<()> { 217 | let document_subdirectory = self.database_subdirectory(); 218 | let removed = self.index.deduplicate()?; 219 | removed.into_par_iter().for_each(|x| { 220 | let _ = std::fs::remove_file(format!("{}/{}.lz4", document_subdirectory, x)); 221 | }); 222 | self.save_database(None)?; 223 | Ok(()) 224 | } 225 | 226 | /// Insert documents into the database.\ 227 | /// Consider batching insertions as inserting too many documents at once may be memory-intensive. 228 | /// 229 | /// # Arguments 230 | /// 231 | /// * `documents` - A vector of documents to be inserted. 232 | pub fn insert_documents(&self, documents: &Vec) -> anyhow::Result<()> { 233 | let new_embeddings: Vec> = self.inner.model.embed_documents(documents)?; 234 | self.insert_records(&new_embeddings, documents) 235 | } 236 | 237 | /// Insert embedding-byte pairs into the database.\ 238 | /// Consider batching insertions as inserting too many records at once may be memory-intensive. 239 | /// 240 | /// # Arguments 241 | /// 242 | /// * `embeddings` - A list of embedding vectors to insert. 243 | /// 244 | /// * `documents` - A list of documents to pair with the embedding vectors. 245 | pub fn insert_records( 246 | &self, 247 | embeddings: &Vec>, 248 | documents: &Vec, 249 | ) -> anyhow::Result<()> { 250 | let embedding_ids = self.index.add(embeddings)?; 251 | self.save_documents_to_disk(&embedding_ids, documents)?; 252 | self.save_database(None)?; 253 | Ok(()) 254 | } 255 | 256 | /// Query records from the database. 257 | /// 258 | /// # Arguments 259 | /// 260 | /// * `documents` - A list of query documents. 261 | /// 262 | /// * `number_of_results` - The maximum number of approximate nearest neighbours to return for each query document. 263 | /// 264 | /// # Returns 265 | /// 266 | /// The records for documents that are most similar to the query documents. 267 | pub fn query_documents( 268 | &self, 269 | documents: &[Bytes], 270 | number_of_results: usize, 271 | ) -> anyhow::Result>>> { 272 | if self.index.no_vectors() { 273 | return Ok(DashMap::new()); 274 | } 275 | let query_embeddings = self.inner.model.embed_documents(documents)?; 276 | self.query_vectors(&query_embeddings, number_of_results) 277 | } 278 | 279 | /// Query records from the database. 280 | /// 281 | /// # Arguments 282 | /// 283 | /// * `vectors` - A list of query vectors. 284 | /// 285 | /// * `number_of_results` - The maximum number of approximate nearest neighbours to return for each query vector. 286 | /// 287 | /// # Returns 288 | /// 289 | /// The records for documents that are most similar to the query vectors. 290 | pub fn query_vectors( 291 | &self, 292 | vectors: &Vec>, 293 | number_of_results: usize, 294 | ) -> anyhow::Result>>> { 295 | if self.index.no_vectors() { 296 | return Ok(DashMap::new()); 297 | } 298 | let results = DashMap::new(); 299 | vectors.into_par_iter().enumerate().for_each(|(idx, x)| { 300 | let mut neighbours = self 301 | .index 302 | .search(x, number_of_results, &self.inner.metric) 303 | .unwrap_or_default(); 304 | neighbours.par_sort_unstable_by_key(|n| n.1); 305 | let neighbour_ids: DashSet<_> = neighbours.into_iter().map(|(id, _)| id).collect(); 306 | results.insert( 307 | idx, 308 | self.read_documents_from_disk(&neighbour_ids) 309 | .unwrap_or_default(), 310 | ); 311 | }); 312 | Ok(results) 313 | } 314 | 315 | /// Save documents to disk. 316 | /// 317 | /// # Arguments 318 | /// 319 | /// * `embedding_ids` - A list of document IDs to be inserted. 320 | /// 321 | /// * `documents` - A list of documents to be inserted. 322 | pub fn save_documents_to_disk( 323 | &self, 324 | embedding_ids: &Vec, 325 | documents: &Vec, 326 | ) -> anyhow::Result<()> { 327 | let document_subdirectory = self.database_subdirectory(); 328 | std::fs::create_dir_all(document_subdirectory.clone())?; 329 | embedding_ids 330 | .par_iter() 331 | .zip(documents.par_iter()) 332 | .map(|(id, document)| -> anyhow::Result<()> { 333 | let mut reader = BufReader::new(Cursor::new(document)); 334 | let file = OpenOptions::new() 335 | .read(true) 336 | .write(true) 337 | .create(true) 338 | .open(format!("{}/{}.lz4", document_subdirectory, id))?; 339 | let buf = BufWriter::new(file); 340 | let mut compressor = lz4_flex::frame::FrameEncoder::new(buf); 341 | io::copy(&mut reader, &mut compressor)?; 342 | compressor.finish()?; 343 | Ok(()) 344 | }) 345 | .collect::>>()?; 346 | Ok(()) 347 | } 348 | 349 | /// Read documents from disk. 350 | /// 351 | /// # Arguments 352 | /// 353 | /// * `documents` - A set of document IDs to be read. 354 | /// 355 | /// # Returns 356 | /// 357 | /// A map of document IDs to the bytes of their corresponding documents. 358 | pub fn read_documents_from_disk( 359 | &self, 360 | documents: &DashSet, 361 | ) -> anyhow::Result>> { 362 | let document_subdirectory = self.database_subdirectory(); 363 | let results = DashMap::new(); 364 | documents 365 | .into_par_iter() 366 | .map(|document_index| -> anyhow::Result<()> { 367 | let file = OpenOptions::new() 368 | .read(true) 369 | .open(format!("{}/{}.lz4", document_subdirectory, *document_index))?; 370 | let buf = BufReader::new(file); 371 | let mut decompressor = lz4_flex::frame::FrameDecoder::new(buf); 372 | let mut writer = BufWriter::new(Vec::new()); 373 | io::copy(&mut decompressor, &mut writer)?; 374 | let document = writer.into_inner()?; 375 | results.insert(*document_index, document); 376 | Ok(()) 377 | }) 378 | .collect::>>()?; 379 | Ok(results) 380 | } 381 | } 382 | -------------------------------------------------------------------------------- /src/database/default/audio.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | database::core::Database, 3 | distance::CosineDistance, 4 | model::{audio::VitBasePatch16_224, core::DIM_VIT_BASE_PATCH16_224}, 5 | }; 6 | 7 | /// The default distance metric for audio embeddings. 8 | pub type DefaultAudioMetric = CosineDistance; 9 | 10 | /// The default embedding model for audio embeddings. 11 | pub type DefaultAudioModel = VitBasePatch16_224; 12 | 13 | /// A database containing sounds and their embeddings. 14 | pub type DefaultAudioDatabase = 15 | Database; 16 | -------------------------------------------------------------------------------- /src/database/default/image.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | database::core::Database, 3 | distance::CosineDistance, 4 | model::{core::DIM_VIT_BASE_PATCH16_224, image::VitBasePatch16_224}, 5 | }; 6 | 7 | /// The default distance metric for image embeddings. 8 | pub type DefaultImageMetric = CosineDistance; 9 | 10 | /// The default embedding model for image embeddings. 11 | pub type DefaultImageModel = VitBasePatch16_224; 12 | 13 | /// A database containing images and their embeddings. 14 | pub type DefaultImageDatabase = 15 | Database; 16 | -------------------------------------------------------------------------------- /src/database/default/mod.rs: -------------------------------------------------------------------------------- 1 | /// Default configuration for an audio database. 2 | pub mod audio; 3 | /// Default configuration for an image database. 4 | pub mod image; 5 | /// Default configuration for a text database. 6 | pub mod text; 7 | -------------------------------------------------------------------------------- /src/database/default/text.rs: -------------------------------------------------------------------------------- 1 | use crate::{ 2 | database::core::Database, 3 | distance::L2SquaredDistance, 4 | model::{core::DIM_BGESMALL_EN_1_5, text::BGESmallEn1_5}, 5 | }; 6 | 7 | /// The default distance metric for text embeddings. 8 | pub type DefaultTextMetric = L2SquaredDistance; 9 | 10 | /// The default embedding model for text embeddings. 11 | pub type DefaultTextModel = BGESmallEn1_5; 12 | 13 | /// A database containing texts and their embeddings. 14 | pub type DefaultTextDatabase = Database; 15 | -------------------------------------------------------------------------------- /src/database/index/lsh.rs: -------------------------------------------------------------------------------- 1 | use std::ops::Deref; 2 | 3 | use crate::{distance::DistanceUnit, Embedding, EmbeddingPrecision, KEYSPACE}; 4 | use dashmap::DashSet; 5 | use fjall::{KvSeparationOptions, PartitionCreateOptions, PartitionHandle, PersistMode}; 6 | use rand::seq::IteratorRandom; 7 | use rayon::{ 8 | iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator}, 9 | slice::ParallelSliceMut, 10 | }; 11 | use serde::{Deserialize, Serialize}; 12 | use simsimd::SpatialSimilarity; 13 | use space::Metric; 14 | use uuid::Uuid; 15 | 16 | #[derive(Debug, Clone, Serialize, Deserialize)] 17 | /// An `N`-dimensional hyperplane; a hyperplane is a generalisation of a line (which has one dimension) or plane (which has two dimensions). 18 | /// 19 | /// It is defined when the dot product of a normal vector and some other vector, plus a constant, equals zero. 20 | pub struct Hyperplane { 21 | /// The vector normal to the hyperplane. 22 | pub coefficients: Embedding, 23 | /// The offset of the hyperplane. 24 | pub constant: EmbeddingPrecision, 25 | } 26 | 27 | impl Hyperplane { 28 | /// Calculates if a point is 'above' the hyperplane. 29 | /// 30 | /// A point is 'above' a hyperplane when it is pointing in the same direction as the hyperplane's normal vector. 31 | /// 32 | /// # Arguments 33 | /// 34 | /// * `point` - The point which may be above, on, or below the hyperplane. 35 | /// 36 | /// # Returns 37 | /// 38 | /// If the given point is above the hyperplane. 39 | pub fn point_is_above(&self, point: &Embedding) -> bool { 40 | EmbeddingPrecision::dot(self.coefficients.deref(), point.deref()).unwrap_or_default() 41 | + self.constant as f64 42 | >= 0.0 43 | } 44 | } 45 | 46 | #[derive(Debug, Clone, Serialize, Deserialize)] 47 | enum Node { 48 | Inner(Box>), 49 | Leaf(Box), 50 | } 51 | 52 | #[derive(Debug, Clone, Serialize, Deserialize)] 53 | struct InnerNode { 54 | hyperplane: Hyperplane, 55 | left_node: Node, 56 | right_node: Node, 57 | } 58 | 59 | #[derive(Debug, Clone, Serialize, Deserialize)] 60 | struct LeafNode(Vec); 61 | 62 | #[derive(Clone)] 63 | struct KeyValue { 64 | keyspace: fjall::Keyspace, 65 | trees: PartitionHandle, 66 | embeddings: PartitionHandle, 67 | } 68 | 69 | impl KeyValue { 70 | fn new(uuid: &Uuid) -> anyhow::Result { 71 | let keyspace = fjall::Config::new(format!("{}-keyspace", uuid.as_simple())).open()?; 72 | let trees = keyspace.open_partition( 73 | &format!("{}-trees", uuid.as_simple()), 74 | PartitionCreateOptions::default().with_kv_separation(KvSeparationOptions::default()), 75 | )?; 76 | let embeddings = keyspace.open_partition( 77 | &format!("{}-embeddings", uuid.as_simple()), 78 | PartitionCreateOptions::default().with_kv_separation(KvSeparationOptions::default()), 79 | )?; 80 | Ok(Self { 81 | keyspace, 82 | trees, 83 | embeddings, 84 | }) 85 | } 86 | 87 | fn save(&self) -> anyhow::Result<()> { 88 | Ok(self.keyspace.persist(PersistMode::SyncAll)?) 89 | } 90 | 91 | fn upsert_embedding(&self, vec_id: &Uuid, embedding: &Embedding) -> anyhow::Result<()> { 92 | self.embeddings.insert( 93 | vec_id.as_bytes(), 94 | bincode::serde::encode_to_vec(embedding, bincode::config::legacy())?, 95 | )?; 96 | self.save() 97 | } 98 | 99 | fn upsert_tree(&self, id: &Uuid, tree: &Node) -> anyhow::Result<()> { 100 | self.trees.insert( 101 | id.as_bytes(), 102 | bincode::serde::encode_to_vec(tree, bincode::config::legacy())?, 103 | )?; 104 | self.save() 105 | } 106 | 107 | fn embedding(&self, idx: &Uuid) -> Embedding { 108 | bincode::serde::decode_from_slice( 109 | &self 110 | .embeddings 111 | .get(idx.as_bytes()) 112 | .ok() 113 | .flatten() 114 | .unwrap_or(fjall::Slice::from(vec![])), 115 | bincode::config::legacy(), 116 | ) 117 | .unwrap_or_default() 118 | .0 119 | } 120 | } 121 | 122 | #[derive(Debug, Clone, Serialize, Deserialize)] 123 | /// Creation options for an [`LSHIndex`]. 124 | pub struct LSHIndexOptions { 125 | /// The maximum number of vectors allowed on one side of a hyperplane; as the maximum node size decreases, the number of partitions in the database grows, increasing query accuracy ('recall') but decreasing performance. 126 | pub max_node_size: usize, 127 | /// The number of trees (ie, the number of root nodes) in the index. 128 | pub num_trees: usize, 129 | } 130 | 131 | impl Default for LSHIndexOptions { 132 | fn default() -> Self { 133 | Self { 134 | max_node_size: 5, 135 | num_trees: 15, 136 | } 137 | } 138 | } 139 | 140 | /// An implementation of [the random projection method of locality sensitive hashing (LSH)](https://en.wikipedia.org/wiki/Locality-sensitive_hashing#Random_projection) as a data structure for use as a database index. 141 | /// 142 | /// This index stores vectors on disk, minimising memory usage. 143 | /// Memory-mapped file IO [is *not* used](https://db.cs.cmu.edu/mmap-cidr2022/). 144 | #[derive(Clone)] 145 | pub struct LSHIndex { 146 | options: LSHIndexOptions, 147 | kv: KeyValue, 148 | } 149 | 150 | impl LSHIndex { 151 | /// Construct a new [LSHIndex]. 152 | /// 153 | /// # Arguments 154 | /// 155 | /// * `uuid` - The UUID of the database. 156 | /// 157 | /// * `options` - The creation options for the index. 158 | /// 159 | /// # Returns 160 | /// 161 | /// An [LSHIndex]. 162 | pub fn new(uuid: &Uuid, options: &LSHIndexOptions) -> anyhow::Result { 163 | Ok(Self { 164 | options: options.clone(), 165 | kv: KeyValue::new(uuid)?, 166 | }) 167 | } 168 | 169 | /// Persist the index to storage. 170 | pub fn save(&self) -> anyhow::Result<()> { 171 | self.kv.save() 172 | } 173 | 174 | fn subtract(lhs: &Embedding, rhs: &Embedding) -> Embedding { 175 | lhs.iter() 176 | .zip(rhs.iter()) 177 | .map(|(a, b)| a - b) 178 | .collect::>() 179 | .try_into() 180 | .unwrap_or_default() 181 | } 182 | 183 | fn average(lhs: &Embedding, rhs: &Embedding) -> Embedding { 184 | lhs.iter() 185 | .zip(rhs.iter()) 186 | .map(|(a, b)| (a + b) / 2.0) 187 | .collect::>() 188 | .try_into() 189 | .unwrap_or_default() 190 | } 191 | 192 | fn build_hyperplane( 193 | &self, 194 | indexes: &Vec, 195 | ) -> anyhow::Result<(Hyperplane, Vec, Vec)> { 196 | // Pick two random vectors 197 | let samples: Vec<_> = self 198 | .kv 199 | .embeddings 200 | .iter() 201 | .choose_multiple(&mut rand::rng(), 2); 202 | 203 | let empty_slice = fjall::Slice::from(vec![]); 204 | let a = samples 205 | .first() 206 | .and_then(|x| x.as_ref().map(|y| &y.1).ok()) 207 | .unwrap_or(&empty_slice); 208 | let b = samples 209 | .get(1) 210 | .and_then(|x| x.as_ref().map(|y| &y.1).ok()) 211 | .unwrap_or(&empty_slice); 212 | let (a, b) = ( 213 | bincode::serde::decode_from_slice(a, bincode::config::legacy()) 214 | .unwrap_or_default() 215 | .0, 216 | bincode::serde::decode_from_slice(b, bincode::config::legacy()) 217 | .unwrap_or_default() 218 | .0, 219 | ); 220 | 221 | // Use the two random points to make a hyperplane orthogonal to a line connecting the two points 222 | let coefficients = Self::subtract(&b, &a); 223 | let point_on_plane = Self::average(&a, &b); 224 | let constant = -EmbeddingPrecision::dot(coefficients.deref(), point_on_plane.deref()) 225 | .unwrap_or_default() as EmbeddingPrecision; 226 | 227 | let hyperplane = Hyperplane { 228 | coefficients, 229 | constant, 230 | }; 231 | 232 | // For each given vector ID, classify a vector as above or below the hyperplane. 233 | let above: DashSet = DashSet::new(); 234 | let below: DashSet = DashSet::new(); 235 | 236 | indexes.into_par_iter().for_each(|id| { 237 | match hyperplane.point_is_above(&self.kv.embedding(id)) { 238 | true => above.insert(*id), 239 | false => below.insert(*id), 240 | }; 241 | }); 242 | 243 | Ok(( 244 | hyperplane, 245 | above.into_par_iter().collect(), 246 | below.into_par_iter().collect(), 247 | )) 248 | } 249 | 250 | fn build_a_tree(&self, indexes: &Vec) -> anyhow::Result> { 251 | match indexes.len() < self.options.max_node_size { 252 | true => Ok(Node::Leaf(Box::new(LeafNode(indexes.clone())))), 253 | false => { 254 | // If there are too many indices to fit into a leaf node, recursively build trees to spread the indices across leaf nodes. 255 | let (hyperplane, above, below) = self.build_hyperplane(indexes)?; 256 | 257 | let node_above = self.build_a_tree(&above)?; 258 | let node_below = self.build_a_tree(&below)?; 259 | 260 | Ok(Node::Inner(Box::new(InnerNode { 261 | hyperplane, 262 | left_node: node_below, 263 | right_node: node_above, 264 | }))) 265 | } 266 | } 267 | } 268 | 269 | /// Remove duplicate embedding vectors from the index. 270 | pub fn deduplicate(&self) -> anyhow::Result> { 271 | let seen: DashSet> = DashSet::new(); 272 | let mut to_remove = Vec::new(); 273 | for kv in self.kv.embeddings.iter() { 274 | let (k, v) = kv?; 275 | let id = Uuid::from_slice(&k)?; 276 | let embedding: Embedding = 277 | bincode::serde::decode_from_slice(&v, bincode::config::legacy())?.0; 278 | // The embedding itself cannot be hashed, so we look at its bits 279 | let embedding_bits: Vec<_> = embedding.iter().map(|x| x.to_bits()).collect(); 280 | match seen.contains(&embedding_bits) { 281 | true => to_remove.push(id), 282 | false => { 283 | seen.insert(embedding_bits); 284 | } 285 | } 286 | } 287 | self.remove(&to_remove) 288 | } 289 | 290 | fn tree_result, Unit = DistanceUnit> + Send + Sync>( 291 | &self, 292 | query: &Embedding, 293 | n: i32, 294 | tree: &Node, 295 | candidates: &DashSet, 296 | metric: &Met, 297 | ) -> anyhow::Result { 298 | match tree { 299 | Node::Leaf(leaf_node) => { 300 | let leaf_values_index = &(leaf_node.0); 301 | match leaf_values_index.len() < n as usize { 302 | true => { 303 | // If there are less than `n` vectors in this leaf node, they're all part of the candidate list. 304 | leaf_values_index.into_par_iter().for_each(|i| { 305 | candidates.insert(*i); 306 | }); 307 | Ok(leaf_values_index.len() as i32) 308 | } 309 | false => { 310 | let mut sorted_candidates = leaf_values_index 311 | .into_par_iter() 312 | .map(|idx| { 313 | let curr_vector: Embedding = self.kv.embedding(idx); 314 | (idx, metric.distance(&curr_vector, query)) 315 | }) 316 | .collect::>(); 317 | sorted_candidates 318 | .par_sort_unstable_by(|a, b| a.1.partial_cmp(&b.1).unwrap()); 319 | let top_candidates: Vec = sorted_candidates 320 | .iter() 321 | .take(n as usize) 322 | .map(|(idx, _)| **idx) 323 | .collect(); 324 | 325 | top_candidates.into_par_iter().for_each(|i| { 326 | candidates.insert(i); 327 | }); 328 | 329 | Ok(n) 330 | } 331 | } 332 | } 333 | Node::Inner(inner_node) => { 334 | let is_above = inner_node.hyperplane.point_is_above(query); 335 | let (main, backup) = match is_above { 336 | true => (&inner_node.right_node, &inner_node.left_node), 337 | false => (&inner_node.left_node, &inner_node.right_node), 338 | }; 339 | 340 | Ok( 341 | match self.tree_result(query, n, main, candidates, metric)? { 342 | k if k < n => self.tree_result(query, n - k, backup, candidates, metric)?, 343 | k => k, 344 | }, 345 | ) 346 | } 347 | } 348 | } 349 | 350 | fn insert( 351 | &self, 352 | current_node: &mut Node, 353 | embedding: &Embedding, 354 | vec_id: Uuid, 355 | ) -> anyhow::Result<()> { 356 | match current_node { 357 | Node::Inner(inner_node) => { 358 | let is_above = inner_node.hyperplane.point_is_above(embedding); 359 | 360 | let next_node = match is_above { 361 | true => &mut inner_node.right_node, 362 | false => &mut inner_node.left_node, 363 | }; 364 | 365 | self.insert(next_node, embedding, vec_id)?; 366 | } 367 | Node::Leaf(leaf_node) => { 368 | match leaf_node.0.len() + 1 > self.options.max_node_size { 369 | false => leaf_node.0.push(vec_id), 370 | true => { 371 | // If adding the vector ID to this leaf node would cause it to be too large, split this node. 372 | let mut new_indexes = leaf_node.0.clone(); 373 | new_indexes.push(vec_id); 374 | 375 | let result_node = self.build_a_tree(&new_indexes)?; 376 | *current_node = result_node; 377 | } 378 | } 379 | } 380 | } 381 | Ok(()) 382 | } 383 | 384 | /// Whether or not the index is empty. 385 | /// 386 | /// # Returns 387 | /// 388 | /// If the index is empty. Note that this is not the same as having no vectors in the index, as deleting all vectors in existing trees does not delete the trees themselves. 389 | pub fn is_empty(&self) -> bool { 390 | self.no_vectors() || self.no_trees() 391 | } 392 | 393 | /// Whether or not there are no vectors in the index. 394 | /// 395 | /// # Returns 396 | /// 397 | /// If there are no vectors in the index. 398 | pub fn no_vectors(&self) -> bool { 399 | self.kv.embeddings.is_empty().unwrap_or(true) 400 | } 401 | 402 | /// Whether or not there are no trees in the index. 403 | /// 404 | /// # Returns 405 | /// 406 | /// If there are no hyperplanes partitioning the vectors in the index. 407 | pub fn no_trees(&self) -> bool { 408 | self.kv.trees.is_empty().unwrap_or(true) 409 | } 410 | 411 | fn build_index(&self, embeddings: &Vec>) -> anyhow::Result> { 412 | let vector_ids = embeddings 413 | .par_iter() 414 | .map(|embedding| { 415 | let vec_id = Uuid::now_v7(); 416 | self.kv.upsert_embedding(&vec_id, embedding)?; 417 | Ok(vec_id) 418 | }) 419 | .collect::>>()?; 420 | (0..self.options.num_trees) 421 | .into_par_iter() 422 | .map(|_| { 423 | let id = Uuid::now_v7(); 424 | let tree = self.build_a_tree(&vector_ids)?; 425 | self.kv.upsert_tree(&id, &tree) 426 | }) 427 | .collect::>>()?; 428 | Ok(vector_ids) 429 | } 430 | 431 | /// Adds vectors to the index. 432 | /// 433 | /// # Arguments 434 | /// 435 | /// * `embeddings` - The embedding vectors to add to the index. 436 | /// 437 | /// # Returns 438 | /// 439 | /// A list of the IDs for the added embedding vectors. 440 | pub fn add(&self, embeddings: &Vec>) -> anyhow::Result> { 441 | if self.no_trees() { 442 | return self.build_index(embeddings); 443 | } 444 | 445 | let vector_ids = embeddings 446 | .par_iter() 447 | .map(|embedding| -> anyhow::Result { 448 | let vec_id = Uuid::now_v7(); 449 | self.kv.upsert_embedding(&vec_id, embedding)?; 450 | 451 | for kv in self.kv.trees.iter() { 452 | let (k, v) = kv?; 453 | let id = Uuid::from_slice(&k)?; 454 | let mut tree: Node = 455 | bincode::serde::decode_from_slice(&v, bincode::config::legacy())?.0; 456 | self.insert(&mut tree, embedding, vec_id)?; 457 | self.kv.upsert_tree(&id, &tree)?; 458 | } 459 | 460 | Ok(vec_id) 461 | }) 462 | .collect::>>()?; 463 | 464 | KEYSPACE.persist(PersistMode::SyncAll)?; 465 | Ok(vector_ids) 466 | } 467 | 468 | /// Removes vectors from the index. 469 | /// 470 | /// # Arguments 471 | /// 472 | /// * `embedding_ids` - The IDs of the vectors to remove. 473 | pub fn remove(&self, embedding_ids: &Vec) -> anyhow::Result> { 474 | let removed = DashSet::new(); 475 | 476 | embedding_ids 477 | .par_iter() 478 | .map(|x| -> anyhow::Result<()> { 479 | self.kv 480 | .trees 481 | .iter() 482 | .try_for_each(|kv| -> anyhow::Result<()> { 483 | let (k, v) = kv?; 484 | let id = Uuid::from_slice(&k)?; 485 | let mut tree: Node = 486 | bincode::serde::decode_from_slice(&v, bincode::config::legacy())?.0; 487 | if let Node::Leaf(leaf) = tree { 488 | let leaf_nodes: Vec = 489 | (*leaf).0.into_par_iter().filter(|y| y != x).collect(); 490 | tree = Node::Leaf(Box::new(LeafNode(leaf_nodes))); 491 | self.kv.upsert_tree(&id, &tree)?; 492 | } 493 | Ok(()) 494 | })?; 495 | self.kv.embeddings.remove(x.as_bytes())?; 496 | removed.insert(*x); 497 | Ok(()) 498 | }) 499 | .collect::>()?; 500 | 501 | KEYSPACE.persist(PersistMode::SyncAll)?; 502 | Ok(removed) 503 | } 504 | 505 | /// Delete the contents of the index. 506 | pub fn clear(&self) -> anyhow::Result<()> { 507 | self.kv 508 | .embeddings 509 | .iter() 510 | .map(|kv| -> anyhow::Result<()> { 511 | let (k, _) = kv?; 512 | self.kv.embeddings.remove(k)?; 513 | Ok(()) 514 | }) 515 | .collect::>>()?; 516 | 517 | self.kv 518 | .trees 519 | .iter() 520 | .map(|kv| -> anyhow::Result<()> { 521 | let (k, _) = kv?; 522 | self.kv.embeddings.remove(k)?; 523 | Ok(()) 524 | }) 525 | .collect::>>()?; 526 | 527 | KEYSPACE.persist(PersistMode::SyncAll)?; 528 | Ok(()) 529 | } 530 | 531 | /// Perform an approximate *k* nearest neighbours search. 532 | /// 533 | /// # Arguments 534 | /// 535 | /// * `query` - The query vector. 536 | /// 537 | /// * `top_k` - The number of approximate neighbours to the query vector to return. 538 | /// 539 | /// * `metric` - The distance metric used to evaluate the distances between the vectors in the index and the query vector. 540 | /// 541 | /// # Returns 542 | /// 543 | /// The IDs of, and distances from, `top_k` approximate nearest neighbours of the query vector. 544 | pub fn search, Unit = DistanceUnit> + Send + Sync>( 545 | &self, 546 | query: &Embedding, 547 | top_k: usize, 548 | metric: &Met, 549 | ) -> anyhow::Result> { 550 | let candidates = DashSet::new(); 551 | 552 | for kv in self.kv.trees.iter() { 553 | let (_, v) = kv?; 554 | let tree: Node = bincode::serde::decode_from_slice(&v, bincode::config::legacy())?.0; 555 | self.tree_result(query, top_k as i32, &tree, &candidates, metric)?; 556 | } 557 | let mut sorted_candidates = candidates 558 | .into_par_iter() 559 | .map(|idx| (idx, metric.distance(&self.kv.embedding(&idx), query))) 560 | .collect::>(); 561 | sorted_candidates.par_sort_unstable_by(|a, b| { 562 | a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal) 563 | }); 564 | Ok(sorted_candidates.into_iter().take(top_k).collect()) 565 | } 566 | } 567 | -------------------------------------------------------------------------------- /src/database/index/mod.rs: -------------------------------------------------------------------------------- 1 | /// Implementation of a locality-sensitive hashing approach to vector database indexing. 2 | /// 3 | /// Derived from [a blog post by Fennel](https://fennel.ai/blog/vector-search-in-200-lines-of-rust/) outlining the approach used by [Spotify's Annoy ('Approximate Nearest Neighbors Oh Yeah')](https://github.com/spotify/annoy). 4 | pub mod lsh; 5 | -------------------------------------------------------------------------------- /src/database/mod.rs: -------------------------------------------------------------------------------- 1 | /// Core implementation of a database. 2 | pub mod core; 3 | #[cfg(feature = "default_db")] 4 | /// Default configurations of databases. 5 | pub mod default; 6 | /// Implementations of database indices. 7 | pub mod index; 8 | -------------------------------------------------------------------------------- /src/distance.rs: -------------------------------------------------------------------------------- 1 | use std::ops::Deref; 2 | 3 | use crate::{Embedding, EmbeddingPrecision}; 4 | use distances::vectors::{ 5 | bray_curtis, canberra, chebyshev, cosine, euclidean, euclidean_sq, hamming, l3_norm, l4_norm, 6 | manhattan, minkowski, minkowski_p, 7 | }; 8 | use serde::{Deserialize, Serialize}; 9 | use simsimd::SpatialSimilarity; 10 | use space::Metric; 11 | 12 | /// The data type representing the distance between two embeddings. 13 | pub type DistanceUnit = u64; 14 | 15 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 16 | /// The cosine distance metric. 17 | pub struct CosineDistance; 18 | 19 | impl Metric> for CosineDistance { 20 | type Unit = DistanceUnit; 21 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 22 | // Use SIMD if vectors are of same length; otherwise, find distance after truncating longer vector so lengths match 23 | EmbeddingPrecision::cosine(a.deref(), b.deref()) 24 | .map(|c| 1.0 - c) 25 | .map(|x| x.to_bits()) 26 | .unwrap_or( 27 | cosine::<_, EmbeddingPrecision>(a.deref(), b.deref()) 28 | .to_bits() 29 | .into(), 30 | ) 31 | } 32 | } 33 | 34 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 35 | /// The L2-squared distance metric. 36 | pub struct L2SquaredDistance; 37 | 38 | impl Metric> for L2SquaredDistance { 39 | type Unit = DistanceUnit; 40 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 41 | EmbeddingPrecision::sqeuclidean(a.deref(), b.deref()) 42 | .map(|x| x.to_bits()) 43 | .unwrap_or( 44 | euclidean_sq::<_, EmbeddingPrecision>(a.deref(), b.deref()) 45 | .to_bits() 46 | .into(), 47 | ) 48 | } 49 | } 50 | 51 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 52 | /// The Chebyshev distance metric. 53 | pub struct ChebyshevDistance; 54 | 55 | impl Metric> for ChebyshevDistance { 56 | type Unit = DistanceUnit; 57 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 58 | let chebyshev_distance = chebyshev(a.deref(), b.deref()); 59 | chebyshev_distance.to_bits().into() 60 | } 61 | } 62 | 63 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 64 | /// The Canberra distance metric. 65 | pub struct CanberraDistance; 66 | 67 | impl Metric> for CanberraDistance { 68 | type Unit = DistanceUnit; 69 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 70 | let canberra_distance: EmbeddingPrecision = canberra(a.deref(), b.deref()); 71 | canberra_distance.to_bits().into() 72 | } 73 | } 74 | 75 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 76 | /// The Bray-Curtis distance metric. 77 | pub struct BrayCurtisDistance; 78 | 79 | impl Metric> for BrayCurtisDistance { 80 | type Unit = DistanceUnit; 81 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 82 | let bray_curtis_distance: EmbeddingPrecision = bray_curtis(a.deref(), b.deref()); 83 | bray_curtis_distance.to_bits().into() 84 | } 85 | } 86 | 87 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 88 | /// The Manhattan distance metric. 89 | pub struct ManhattanDistance; 90 | 91 | impl Metric> for ManhattanDistance { 92 | type Unit = DistanceUnit; 93 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 94 | let manhattan_distance: EmbeddingPrecision = manhattan(a.deref(), b.deref()); 95 | manhattan_distance.to_bits().into() 96 | } 97 | } 98 | 99 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 100 | /// The L2 distance metric. 101 | pub struct L2Distance; 102 | 103 | impl Metric> for L2Distance { 104 | type Unit = DistanceUnit; 105 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 106 | EmbeddingPrecision::euclidean(a.deref(), b.deref()) 107 | .map(|x| x.to_bits()) 108 | .unwrap_or( 109 | euclidean::<_, EmbeddingPrecision>(a.deref(), b.deref()) 110 | .to_bits() 111 | .into(), 112 | ) 113 | } 114 | } 115 | 116 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 117 | /// The L3 distance metric. 118 | pub struct L3Distance; 119 | 120 | impl Metric> for L3Distance { 121 | type Unit = DistanceUnit; 122 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 123 | let l3_distance: EmbeddingPrecision = l3_norm(a.deref(), b.deref()); 124 | l3_distance.to_bits().into() 125 | } 126 | } 127 | 128 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 129 | /// The L4 distance metric. 130 | pub struct L4Distance; 131 | 132 | impl Metric> for L4Distance { 133 | type Unit = DistanceUnit; 134 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 135 | let l4_distance: EmbeddingPrecision = l4_norm(a.deref(), b.deref()); 136 | l4_distance.to_bits().into() 137 | } 138 | } 139 | 140 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 141 | /// The Hamming distance metric. 142 | pub struct HammingDistance; 143 | 144 | impl Metric> for HammingDistance { 145 | type Unit = DistanceUnit; 146 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 147 | let a_to_bits: Vec = a.iter().map(|x| x.to_bits() as u8).collect(); 148 | let b_to_bits: Vec = b.iter().map(|x| x.to_bits() as u8).collect(); 149 | match a.len() == b.len() { 150 | true => hamming_bitwise_fast::hamming_bitwise_fast( 151 | a_to_bits.as_slice(), 152 | b_to_bits.as_slice(), 153 | ) 154 | .into(), 155 | false => hamming::<_, u32>(a_to_bits.as_slice(), b_to_bits.as_slice()).into(), 156 | } 157 | } 158 | } 159 | 160 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 161 | /// The Minkowski distance metric. 162 | pub struct MinkowskiDistance { 163 | /// The power of the Minkowski distance. 164 | pub power: i32, 165 | } 166 | 167 | impl Metric> for MinkowskiDistance { 168 | type Unit = DistanceUnit; 169 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 170 | let metric = minkowski(self.power); 171 | let distance: EmbeddingPrecision = metric(a.deref(), b.deref()); 172 | distance.to_bits().into() 173 | } 174 | } 175 | 176 | #[derive(Default, Debug, Clone, Serialize, Deserialize)] 177 | /// The p-norm distance metric. 178 | pub struct PNormDistance { 179 | /// The power of the distance metric. 180 | pub power: i32, 181 | } 182 | 183 | impl Metric> for PNormDistance { 184 | type Unit = DistanceUnit; 185 | fn distance(&self, a: &Embedding, b: &Embedding) -> Self::Unit { 186 | let metric = minkowski_p(self.power); 187 | let distance: EmbeddingPrecision = metric(a.deref(), b.deref()); 188 | distance.to_bits().into() 189 | } 190 | } 191 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #![doc = include_str!("../README.md")] 2 | // #![feature(doc_auto_cfg)] 3 | #![warn(missing_docs)] 4 | 5 | /// Implementation of database operations. 6 | pub mod database; 7 | /// Implementation of distance metrics. 8 | pub mod distance; 9 | /// Interface for embedding models. 10 | pub mod model; 11 | 12 | use serde::{Deserialize, Serialize}; 13 | use serde_with::serde_as; 14 | 15 | /// An embedding vector. 16 | #[serde_as] 17 | #[derive(Debug, Clone, Serialize, Deserialize)] 18 | pub struct Embedding(#[serde_as(as = "[_; N]")] [EmbeddingPrecision; N]); 19 | impl std::ops::Deref for Embedding { 20 | type Target = [EmbeddingPrecision; N]; 21 | 22 | fn deref(&self) -> &[EmbeddingPrecision; N] { 23 | &self.0 24 | } 25 | } 26 | impl std::ops::DerefMut for Embedding { 27 | fn deref_mut(&mut self) -> &mut [EmbeddingPrecision; N] { 28 | &mut self.0 29 | } 30 | } 31 | impl Default for Embedding { 32 | fn default() -> Self { 33 | Self([0.0; N]) 34 | } 35 | } 36 | impl From<[EmbeddingPrecision; N]> for Embedding { 37 | fn from(value: [EmbeddingPrecision; N]) -> Self { 38 | Self(value) 39 | } 40 | } 41 | impl TryFrom> for Embedding { 42 | type Error = Vec; 43 | fn try_from(value: Vec) -> Result { 44 | Ok(Self(value.try_into()?)) 45 | } 46 | } 47 | /// The floating-point precision used to represent embedding vectors. 48 | pub type EmbeddingPrecision = f32; 49 | static KEYSPACE: std::sync::LazyLock = std::sync::LazyLock::new(|| { 50 | fjall::Config::new("keyspace") 51 | .open() 52 | .expect("Keyspace should be accessible from disk") 53 | }); 54 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | use bytes::Bytes; 2 | use clap::{command, Parser, Subcommand}; 3 | use indicatif::HumanCount; 4 | use indicatif::ProgressStyle; 5 | use indicatif::{ProgressBar, ProgressDrawTarget}; 6 | use pretty_duration::pretty_duration; 7 | use rayon::iter::IndexedParallelIterator; 8 | use rayon::iter::IntoParallelIterator; 9 | use rayon::iter::IntoParallelRefIterator; 10 | use rayon::iter::ParallelIterator; 11 | use rodio::{Decoder, OutputStream, Sink}; 12 | use serde::Deserialize; 13 | use serde::Serialize; 14 | use space::Metric; 15 | use std::io::BufReader; 16 | use std::io::Cursor; 17 | use std::io::Write; 18 | use std::io::{stdout, BufWriter}; 19 | use std::path::PathBuf; 20 | use ticky::Stopwatch; 21 | use zebra::database::core::Database; 22 | use zebra::database::default::audio::DefaultAudioDatabase; 23 | use zebra::database::default::image::DefaultImageDatabase; 24 | use zebra::database::default::text::DefaultTextDatabase; 25 | use zebra::distance::DistanceUnit; 26 | use zebra::model::core::DatabaseEmbeddingModel; 27 | use zebra::model::core::DIM_BGESMALL_EN_1_5; 28 | use zebra::Embedding; 29 | 30 | #[derive(Parser)] 31 | #[command(version, about, long_about = None, arg_required_else_help(true))] 32 | struct Cli { 33 | #[structopt(subcommand)] 34 | commands: Commands, 35 | #[arg(short, long, global = true)] 36 | database_path: String, 37 | } 38 | 39 | #[derive(Subcommand)] 40 | enum Commands { 41 | #[clap(about = "Text commands.")] 42 | Text(Text), 43 | #[clap(about = "Image commands.")] 44 | Image(Image), 45 | #[clap(about = "Audio commands.")] 46 | Audio(Audio), 47 | } 48 | 49 | #[derive(Parser)] 50 | struct Text { 51 | #[structopt(subcommand)] 52 | text_commands: TextCommands, 53 | } 54 | 55 | #[derive(Parser)] 56 | struct Image { 57 | #[structopt(subcommand)] 58 | image_commands: ImageCommands, 59 | } 60 | 61 | #[derive(Parser)] 62 | struct Audio { 63 | #[structopt(subcommand)] 64 | audio_commands: AudioCommands, 65 | } 66 | 67 | #[derive(Subcommand)] 68 | enum TextCommands { 69 | #[command( 70 | about = "Insert texts into the database.", 71 | arg_required_else_help(true) 72 | )] 73 | Insert { texts: Vec }, 74 | #[command( 75 | about = "Insert texts into the database from files on disk.", 76 | arg_required_else_help(true) 77 | )] 78 | InsertFromFiles { 79 | file_paths: Vec, 80 | #[arg(default_value_t = 100)] 81 | batch_size: usize, 82 | }, 83 | #[command(about = "Query texts from the database.", arg_required_else_help(true))] 84 | Query { 85 | texts: Vec, 86 | #[arg(default_value_t = 1)] 87 | number_of_results: usize, 88 | }, 89 | #[command(about = "Clear the database.")] 90 | Clear, 91 | } 92 | 93 | #[derive(Subcommand)] 94 | enum ImageCommands { 95 | #[command( 96 | about = "Insert images into the database.", 97 | arg_required_else_help(true) 98 | )] 99 | Insert { 100 | file_paths: Vec, 101 | #[arg(default_value_t = 100)] 102 | batch_size: usize, 103 | }, 104 | #[command( 105 | about = "Query images from the database.", 106 | arg_required_else_help(true) 107 | )] 108 | Query { 109 | image_path: PathBuf, 110 | #[arg(default_value_t = 1)] 111 | number_of_results: usize, 112 | }, 113 | #[command(about = "Clear the database.")] 114 | Clear, 115 | } 116 | 117 | #[derive(Subcommand)] 118 | enum AudioCommands { 119 | #[command( 120 | about = "Insert sounds into the database.", 121 | arg_required_else_help(true) 122 | )] 123 | Insert { 124 | file_paths: Vec, 125 | #[arg(default_value_t = 100)] 126 | batch_size: usize, 127 | }, 128 | #[command( 129 | about = "Query sounds from the database.", 130 | arg_required_else_help(true) 131 | )] 132 | Query { 133 | audio_path: PathBuf, 134 | #[arg(default_value_t = 1)] 135 | number_of_results: usize, 136 | }, 137 | #[command(about = "Clear the database.")] 138 | Clear, 139 | } 140 | 141 | fn main() -> anyhow::Result<()> { 142 | let cli = Cli::parse(); 143 | match cli.commands { 144 | Commands::Text(text) => match text.text_commands { 145 | TextCommands::Insert { texts } => { 146 | let mut sw = Stopwatch::start_new(); 147 | let db = 148 | DefaultTextDatabase::open_or_create(&cli.database_path, &Default::default())?; 149 | let mut buffer = BufWriter::new(stdout().lock()); 150 | writeln!(buffer, "Inserting {} text(s).", texts.len())?; 151 | let texts_bytes: Vec<_> = texts.into_par_iter().map(Bytes::from).collect(); 152 | db.insert_documents(&texts_bytes)?; 153 | sw.stop(); 154 | writeln!( 155 | buffer, 156 | "{} embeddings of {} dimensions inserted into the database in {}.", 157 | HumanCount(texts_bytes.len() as u64), 158 | HumanCount(DIM_BGESMALL_EN_1_5 as u64), 159 | pretty_duration(&sw.elapsed(), None) 160 | )?; 161 | } 162 | TextCommands::InsertFromFiles { 163 | file_paths, 164 | batch_size, 165 | } => { 166 | let db = 167 | DefaultTextDatabase::open_or_create(&cli.database_path, &Default::default())?; 168 | insert_from_files(&db, file_paths, batch_size)?; 169 | } 170 | TextCommands::Query { 171 | texts, 172 | number_of_results, 173 | } => { 174 | let mut sw = Stopwatch::start_new(); 175 | let db = 176 | DefaultTextDatabase::open_or_create(&cli.database_path, &Default::default())?; 177 | let mut buffer = BufWriter::new(stdout().lock()); 178 | let num_texts = texts.len(); 179 | writeln!(buffer, "Querying {} text(s).", num_texts)?; 180 | let texts_bytes: Vec<_> = texts.into_par_iter().map(Bytes::from).collect(); 181 | let query_results = db.query_documents(&texts_bytes, number_of_results)?; 182 | sw.stop(); 183 | writeln!( 184 | buffer, 185 | "Queried {} text(s) in {}.", 186 | num_texts, 187 | pretty_duration(&sw.elapsed(), None) 188 | )?; 189 | writeln!(buffer, "Results:")?; 190 | for (idx, result) in query_results { 191 | writeln!(buffer, "{idx}:\n")?; 192 | writeln!(buffer, "{result:#?}")?; 193 | } 194 | } 195 | TextCommands::Clear => { 196 | DefaultTextDatabase::open_or_create(&cli.database_path, &Default::default())? 197 | .clear_database(); 198 | } 199 | }, 200 | Commands::Image(image) => match image.image_commands { 201 | ImageCommands::Insert { 202 | file_paths, 203 | batch_size, 204 | } => { 205 | let db = 206 | DefaultImageDatabase::open_or_create(&cli.database_path, &Default::default())?; 207 | insert_from_files(&db, file_paths, batch_size)?; 208 | } 209 | ImageCommands::Query { 210 | image_path, 211 | number_of_results, 212 | } => { 213 | let mut sw = Stopwatch::start_new(); 214 | let db = 215 | DefaultImageDatabase::open_or_create(&cli.database_path, &Default::default())?; 216 | let mut buffer = BufWriter::new(stdout().lock()); 217 | let image_print_config = viuer::Config { 218 | transparent: true, 219 | premultiplied_alpha: false, 220 | absolute_offset: false, 221 | x: 0, 222 | y: 0, 223 | restore_cursor: true, 224 | width: None, 225 | height: None, 226 | truecolor: true, 227 | use_kitty: true, 228 | use_iterm: true, 229 | #[cfg(feature = "sixel")] 230 | use_sixel: true, 231 | }; 232 | writeln!(buffer, "Querying image.")?; 233 | let image_bytes = std::fs::read(image_path).unwrap_or_default().into(); 234 | let query_results = db.query_documents(&[image_bytes], number_of_results)?; 235 | sw.stop(); 236 | writeln!( 237 | buffer, 238 | "Queried image in {}.", 239 | pretty_duration(&sw.elapsed(), None) 240 | )?; 241 | writeln!(buffer, "Results:")?; 242 | for (idx, result) in query_results { 243 | writeln!(buffer, "{idx}:\n")?; 244 | for (_, image) in result { 245 | let img = image::load_from_memory(&image)?; 246 | let _ = viuer::print(&img, &image_print_config); 247 | } 248 | } 249 | } 250 | ImageCommands::Clear => { 251 | DefaultImageDatabase::open_or_create(&cli.database_path, &Default::default())? 252 | .clear_database(); 253 | } 254 | }, 255 | Commands::Audio(audio) => match audio.audio_commands { 256 | AudioCommands::Insert { 257 | file_paths, 258 | batch_size, 259 | } => { 260 | let db = 261 | DefaultAudioDatabase::open_or_create(&cli.database_path, &Default::default())?; 262 | insert_from_files(&db, file_paths, batch_size)?; 263 | } 264 | AudioCommands::Query { 265 | audio_path, 266 | number_of_results, 267 | } => { 268 | let mut sw = Stopwatch::start_new(); 269 | let db = 270 | DefaultAudioDatabase::open_or_create(&cli.database_path, &Default::default())?; 271 | let (_stream, stream_handle) = OutputStream::try_default()?; 272 | let sink = Sink::try_new(&stream_handle)?; 273 | let mut buffer = BufWriter::new(stdout().lock()); 274 | writeln!(buffer, "Querying sound.")?; 275 | let audio_bytes = std::fs::read(audio_path).unwrap_or_default().into(); 276 | let query_results = db.query_documents(&[audio_bytes], number_of_results)?; 277 | sw.stop(); 278 | writeln!( 279 | buffer, 280 | "Queried sound in {}.", 281 | pretty_duration(&sw.elapsed(), None) 282 | )?; 283 | writeln!(buffer, "Results:")?; 284 | for (idx, result) in query_results { 285 | writeln!(buffer, "{idx}:\n")?; 286 | for (id, audio) in result { 287 | writeln!(buffer, "Playing {} … ", id.simple())?; 288 | let reader = BufReader::new(Cursor::new(audio)); 289 | let source = Decoder::new(reader)?; 290 | sink.append(source); 291 | sink.sleep_until_end(); 292 | } 293 | } 294 | } 295 | AudioCommands::Clear => { 296 | DefaultAudioDatabase::open_or_create(&cli.database_path, &Default::default())? 297 | .clear_database(); 298 | } 299 | }, 300 | } 301 | Ok(()) 302 | } 303 | 304 | fn progress_bar_style() -> anyhow::Result { 305 | Ok(ProgressStyle::with_template("[{elapsed} elapsed, {eta} remaining ({duration} total)] {wide_bar:.cyan/blue} {human_pos} of {human_len} ({percent}%) {msg}")?) 306 | } 307 | 308 | fn insert_from_files( 309 | db: &Database, 310 | file_paths: Vec, 311 | batch_size: usize, 312 | ) -> anyhow::Result<()> 313 | where 314 | for<'de> Met: Metric, Unit = DistanceUnit> 315 | + Default 316 | + Serialize 317 | + Send 318 | + Sync 319 | + Deserialize<'de>, 320 | for<'de> Mod: DatabaseEmbeddingModel + Default + Serialize + Send + Sync + Deserialize<'de>, 321 | { 322 | let mut sw = Stopwatch::start_new(); 323 | let num_documents = file_paths.len(); 324 | println!( 325 | "Inserting documents from {} file(s).", 326 | HumanCount(num_documents as u64) 327 | ); 328 | let progress_bar = 329 | ProgressBar::with_draw_target(Some(num_documents as u64), ProgressDrawTarget::hidden()); 330 | progress_bar.set_style(progress_bar_style()?); 331 | let documents: Vec<_> = file_paths 332 | .par_iter() 333 | .filter_map(|x| std::fs::read(x).ok().map(|y| y.into())) 334 | .collect(); 335 | documents 336 | .into_par_iter() 337 | .chunks(batch_size) 338 | .map(|document_batch| -> anyhow::Result<()> { 339 | let mut batch_sw = Stopwatch::start_new(); 340 | db.insert_documents(&document_batch)?; 341 | batch_sw.stop(); 342 | progress_bar.println(format!( 343 | "{} embeddings of {} dimensions inserted into the database in {}.", 344 | HumanCount(document_batch.len() as u64), 345 | HumanCount(N as u64), 346 | pretty_duration(&batch_sw.elapsed(), None) 347 | )); 348 | progress_bar.inc(batch_size as u64); 349 | if progress_bar.is_hidden() { 350 | progress_bar.set_draw_target(ProgressDrawTarget::stderr_with_hz(100)); 351 | } 352 | Ok(()) 353 | }) 354 | .collect::>>()?; 355 | sw.stop(); 356 | progress_bar.println(format!( 357 | "Inserted {} document(s) in {}.", 358 | num_documents, 359 | pretty_duration(&sw.elapsed(), None) 360 | )); 361 | Ok(()) 362 | } 363 | -------------------------------------------------------------------------------- /src/model/audio.rs: -------------------------------------------------------------------------------- 1 | use super::{ 2 | core::{DatabaseEmbeddingModel, DIM_VIT_BASE_PATCH16_224}, 3 | image::ImageEmbeddingModel, 4 | }; 5 | use crate::{Embedding, EmbeddingPrecision}; 6 | use anyhow::anyhow; 7 | use bytes::Bytes; 8 | use candle_core::{DType, Tensor}; 9 | use candle_nn::VarBuilder; 10 | use candle_transformers::models::vit; 11 | use rayon::iter::{IntoParallelIterator, ParallelIterator}; 12 | use serde::{Deserialize, Serialize}; 13 | use sonogram::{ColourGradient, FrequencyScale, SpecOptionsBuilder}; 14 | use std::io::Cursor; 15 | use symphonia::core::{ 16 | audio::Signal, 17 | codecs::{DecoderOptions, CODEC_TYPE_NULL}, 18 | formats::FormatOptions, 19 | io::MediaSourceStream, 20 | meta::MetadataOptions, 21 | probe::Hint, 22 | }; 23 | 24 | /// Interface provided by audio embedding models; these models are a subset of image embedding models. 25 | pub trait AudioEmbeddingModel: ImageEmbeddingModel { 26 | /// Decodes the samples of an audio file. 27 | /// 28 | /// # Arguments 29 | /// 30 | /// * `audio` - The raw bytes of an audio file. 31 | /// 32 | /// # Returns 33 | /// 34 | /// An `i16` vector of decoded samples, and the sample rate of the audio. 35 | fn audio_to_data(&self, audio: &Bytes) -> anyhow::Result<(Vec, u32)> { 36 | let mss = MediaSourceStream::new(Box::new(Cursor::new(audio.to_vec())), Default::default()); 37 | let meta_opts: MetadataOptions = Default::default(); 38 | let fmt_opts: FormatOptions = Default::default(); 39 | let probed = 40 | symphonia::default::get_probe().format(&Hint::new(), mss, &fmt_opts, &meta_opts)?; 41 | let mut format = probed.format; 42 | let track = format 43 | .tracks() 44 | .into_par_iter() 45 | .find_any(|t| t.codec_params.codec != CODEC_TYPE_NULL) 46 | .ok_or(anyhow!("No tracks found in audio … "))?; 47 | let dec_opts: DecoderOptions = Default::default(); 48 | let mut decoder = symphonia::default::get_codecs().make(&track.codec_params, &dec_opts)?; 49 | let track_id = track.id; 50 | let mut sample_rate = 0; 51 | let mut data = Vec::new(); 52 | 53 | while let Ok(packet) = format.next_packet() { 54 | while !format.metadata().is_latest() { 55 | format.metadata().pop(); 56 | } 57 | if packet.track_id() != track_id { 58 | continue; 59 | } 60 | match decoder.decode(&packet) { 61 | Ok(decoded) => { 62 | let decoded = decoded.make_equivalent::(); 63 | sample_rate = decoded.spec().rate; 64 | let number_channels = decoded.spec().channels.count(); 65 | for i in 0..number_channels { 66 | let samples = decoded.chan(i); 67 | data.extend_from_slice(samples); 68 | } 69 | } 70 | Err(_) => continue, 71 | } 72 | } 73 | 74 | Ok((data, sample_rate)) 75 | } 76 | 77 | /// Convert an audio file into a logarithm-scale spectrogram for use with image embedding models. 78 | /// 79 | /// # Arguments 80 | /// 81 | /// `audio` - The raw bytes of an audio file. 82 | /// 83 | /// # Returns 84 | /// 85 | /// A spectrogram of the audio as an ImageNet-normalised tensor with shape [3 224 224]. 86 | fn audio_to_image_tensor224(&self, audio: &Bytes) -> anyhow::Result { 87 | let (data, sample_rate) = self.audio_to_data(audio)?; 88 | let mut spectrograph = SpecOptionsBuilder::new(512) 89 | .load_data_from_memory(data, sample_rate) 90 | .normalise() 91 | .build() 92 | .ok() 93 | .ok_or(anyhow!("Unable to compute spectrograph … "))?; 94 | let mut spectrogram = spectrograph.compute(); 95 | let mut gradient = ColourGradient::rainbow_theme(); 96 | let png_bytes = 97 | spectrogram.to_png_in_memory(FrequencyScale::Log, &mut gradient, 224, 224)?; 98 | self.load_image224(&Bytes::from(png_bytes)) 99 | } 100 | } 101 | 102 | /// A model for embedding audio. 103 | #[derive( 104 | Default, Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, 105 | )] 106 | pub struct VitBasePatch16_224; 107 | impl ImageEmbeddingModel for VitBasePatch16_224 {} 108 | impl AudioEmbeddingModel for VitBasePatch16_224 {} 109 | 110 | impl DatabaseEmbeddingModel for VitBasePatch16_224 { 111 | fn embed_documents( 112 | &self, 113 | documents: &[bytes::Bytes], 114 | ) -> anyhow::Result>> { 115 | let mut result = Vec::new(); 116 | let device = candle_examples::device(false)?; 117 | let api = hf_hub::api::sync::Api::new()?; 118 | let api = api.model("google/vit-base-patch16-224".into()); 119 | let model_file = api.get("model.safetensors")?; 120 | let varbuilder = 121 | unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, &device)? }; 122 | let model = vit::Embeddings::new( 123 | &vit::Config::vit_base_patch16_224(), 124 | false, 125 | varbuilder.pp("vit").pp("embeddings"), 126 | )?; 127 | for document in documents { 128 | let image = self 129 | .audio_to_image_tensor224(document)? 130 | .to_device(&device)?; 131 | let embedding_tensors = model.forward(&image.unsqueeze(0)?, None, false)?; 132 | let embedding_vector = embedding_tensors 133 | .flatten_all()? 134 | .to_vec1::()?; 135 | result.push(embedding_vector); 136 | } 137 | Ok(result 138 | .into_par_iter() 139 | .map(|x| x.try_into().unwrap_or_default()) 140 | .collect()) 141 | } 142 | } 143 | -------------------------------------------------------------------------------- /src/model/core.rs: -------------------------------------------------------------------------------- 1 | use crate::Embedding; 2 | use bytes::Bytes; 3 | use serde::Serialize; 4 | 5 | /// Dimensionality of embeddings produced by the [crate::model::text::BGESmallEn1_5] model. 6 | pub const DIM_BGESMALL_EN_1_5: usize = 384; 7 | 8 | /// Dimensionality of embeddings produced by the [crate::model::image::VitBasePatch16_224] model. 9 | pub const DIM_VIT_BASE_PATCH16_224: usize = 768; 10 | 11 | /// A trait for embedding models that can be used with the database. 12 | pub trait DatabaseEmbeddingModel: Serialize { 13 | /// Embed a vector of documents. 14 | /// 15 | /// # Arguments 16 | /// 17 | /// * `documents` - A vector of documents to be embedded. 18 | /// 19 | /// # Returns 20 | /// 21 | /// A vector of embeddings. 22 | fn embed_documents(&self, documents: &[Bytes]) -> anyhow::Result>>; 23 | 24 | /// Embed a single document. 25 | /// 26 | /// # Arguments 27 | /// 28 | /// * `document` – A single document to be embedded. 29 | /// 30 | /// # Returns 31 | /// 32 | /// An embedding vector. 33 | fn embed(&self, document: Bytes) -> anyhow::Result> { 34 | self.embed_documents(&[document]) 35 | .map(|x| x.into_iter().next().unwrap_or_default()) 36 | } 37 | } 38 | -------------------------------------------------------------------------------- /src/model/image.rs: -------------------------------------------------------------------------------- 1 | use super::core::{DatabaseEmbeddingModel, DIM_VIT_BASE_PATCH16_224}; 2 | use crate::{Embedding, EmbeddingPrecision}; 3 | use bytes::Bytes; 4 | use candle_core::{DType, Tensor}; 5 | use candle_examples::imagenet::{IMAGENET_MEAN, IMAGENET_STD}; 6 | use candle_nn::VarBuilder; 7 | use candle_transformers::models::vit; 8 | use image::ImageReader; 9 | use rayon::iter::{IntoParallelIterator, ParallelIterator}; 10 | use serde::{Deserialize, Serialize}; 11 | use std::io::Cursor; 12 | 13 | /// Interface provided by image embedding models. 14 | pub trait ImageEmbeddingModel { 15 | /// Loads an image from raw bytes with ImageNet normalisation applied, returning a tensor with the shape [3 224 224]. 16 | /// 17 | /// # Arguments 18 | /// 19 | /// * `bytes` - The raw bytes of an image. 20 | /// 21 | /// # Returns 22 | /// 23 | /// A tensor with the shape [3 224 224]; ImageNet normalisation is applied. 24 | fn load_image224(&self, bytes: &Bytes) -> anyhow::Result { 25 | let res = 224_usize; 26 | let img = ImageReader::new(Cursor::new(bytes)) 27 | .with_guessed_format()? 28 | .decode()? 29 | .resize_to_fill( 30 | res as u32, 31 | res as u32, 32 | image::imageops::FilterType::Triangle, 33 | ) 34 | .to_rgb8(); 35 | let data = img.into_raw(); 36 | let data = 37 | Tensor::from_vec(data, (res, res, 3), &candle_core::Device::Cpu)?.permute((2, 0, 1))?; 38 | let mean = Tensor::new(&IMAGENET_MEAN, &candle_core::Device::Cpu)?.reshape((3, 1, 1))?; 39 | let std = Tensor::new(&IMAGENET_STD, &candle_core::Device::Cpu)?.reshape((3, 1, 1))?; 40 | Ok((data.to_dtype(candle_core::DType::F32)? / 255.)? 41 | .broadcast_sub(&mean)? 42 | .broadcast_div(&std)?) 43 | } 44 | } 45 | 46 | /// A model for embedding images. 47 | #[derive( 48 | Default, Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, 49 | )] 50 | pub struct VitBasePatch16_224; 51 | impl ImageEmbeddingModel for VitBasePatch16_224 {} 52 | 53 | impl DatabaseEmbeddingModel for VitBasePatch16_224 { 54 | fn embed_documents( 55 | &self, 56 | documents: &[bytes::Bytes], 57 | ) -> anyhow::Result>> { 58 | let mut result = Vec::new(); 59 | let device = candle_examples::device(false)?; 60 | let api = hf_hub::api::sync::Api::new()?; 61 | let api = api.model("google/vit-base-patch16-224".into()); 62 | let model_file = api.get("model.safetensors")?; 63 | let varbuilder = 64 | unsafe { VarBuilder::from_mmaped_safetensors(&[model_file], DType::F32, &device)? }; 65 | let model = vit::Embeddings::new( 66 | &vit::Config::vit_base_patch16_224(), 67 | false, 68 | varbuilder.pp("vit").pp("embeddings"), 69 | )?; 70 | for document in documents { 71 | let image = self.load_image224(document)?.to_device(&device)?; 72 | let embedding_tensors = model.forward(&image.unsqueeze(0)?, None, false)?; 73 | let embedding_vector = embedding_tensors 74 | .flatten_all()? 75 | .to_vec1::()?; 76 | result.push(embedding_vector); 77 | } 78 | Ok(result 79 | .into_par_iter() 80 | .map(|x| x.try_into().unwrap_or_default()) 81 | .collect()) 82 | } 83 | } 84 | -------------------------------------------------------------------------------- /src/model/mod.rs: -------------------------------------------------------------------------------- 1 | /// Audio embedding models. 2 | pub mod audio; 3 | /// Core model interface. 4 | pub mod core; 5 | /// Image embedding models. 6 | pub mod image; 7 | /// Text embedding models. 8 | pub mod text; 9 | -------------------------------------------------------------------------------- /src/model/text.rs: -------------------------------------------------------------------------------- 1 | use super::core::{DatabaseEmbeddingModel, DIM_BGESMALL_EN_1_5}; 2 | use crate::Embedding; 3 | use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; 4 | use rayon::iter::{IntoParallelIterator, ParallelIterator}; 5 | use serde::{Deserialize, Serialize}; 6 | 7 | /// A model for embedding text. 8 | #[derive( 9 | Default, Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, 10 | )] 11 | pub struct BGESmallEn1_5; 12 | 13 | impl DatabaseEmbeddingModel for BGESmallEn1_5 { 14 | fn embed_documents(&self, documents: &[bytes::Bytes]) -> anyhow::Result>> { 15 | let model = TextEmbedding::try_new( 16 | InitOptions::new(EmbeddingModel::BGESmallENV15).with_show_download_progress(false), 17 | )?; 18 | let embeddings = model.embed( 19 | documents 20 | .into_par_iter() 21 | .map(|x| x.to_vec()) 22 | .filter_map(|x| String::from_utf8(x).ok()) 23 | .collect(), 24 | None, 25 | )?; 26 | Ok(embeddings 27 | .into_par_iter() 28 | .map(|x| x.try_into().unwrap_or_default()) 29 | .collect()) 30 | } 31 | } 32 | --------------------------------------------------------------------------------