diff --git a/Cargo.lock b/Cargo.lock index 4f7f1668..af0e063d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,189 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "415ed64958754dbe991900f3940677e6a7eefb4d7367afd70d642677b0c7d19d" +[[package]] +name = "actix-codec" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f7b0a21988c1bf877cf4759ef5ddaac04c1c9fe808c9142ecb78ba97d97a28a" +dependencies = [ + "bitflags 2.9.1", + "bytes", + "futures-core", + "futures-sink", + "memchr", + "pin-project-lite", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "actix-http" +version = "3.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7926860314cbe2fb5d1f13731e387ab43bd32bca224e82e6e2db85de0a3dba49" +dependencies = [ + "actix-codec", + "actix-rt", + "actix-service", + "actix-utils", + "base64 0.22.1", + "bitflags 2.9.1", + "brotli", + "bytes", + "bytestring", + "derive_more 2.0.1", + "encoding_rs", + "flate2", + "foldhash", + "futures-core", + "h2", + "http 0.2.12", + "httparse", + "httpdate", + "itoa", + "language-tags", + "local-channel", + "mime", + "percent-encoding", + "pin-project-lite", + "rand 0.9.1", + "sha1", + "smallvec 1.15.1", + "tokio", + "tokio-util", + "tracing", + "zstd", +] + +[[package]] +name = "actix-macros" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01ed3140b2f8d422c68afa1ed2e85d996ea619c988ac834d255db32138655cb" +dependencies = [ + "quote", + "syn 2.0.102", +] + +[[package]] +name = "actix-router" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d324164c51f63867b57e73ba5936ea151b8a41a1d23d1031eeb9f70d0236f8" +dependencies = [ + "bytestring", + "cfg-if", + "http 0.2.12", + "regex", + "regex-lite", + "serde", + "tracing", +] + +[[package]] +name = "actix-rt" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92589714878ca59a7626ea19734f0e07a6a875197eec751bb5d3f99e64998c63" +dependencies = [ + "futures-core", + "tokio", +] + +[[package]] +name = "actix-server" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a65064ea4a457eaf07f2fba30b4c695bf43b721790e9530d26cb6f9019ff7502" +dependencies = [ + "actix-rt", + "actix-service", + "actix-utils", + "futures-core", + "futures-util", + "mio", + "socket2", + "tokio", + "tracing", +] + +[[package]] +name = "actix-service" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e46f36bf0e5af44bdc4bdb36fbbd421aa98c79a9bce724e1edeb3894e10dc7f" +dependencies = [ + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "actix-utils" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88a1dcdff1466e3c2488e1cb5c36a71822750ad43839937f85d2f4d9f8b705d8" +dependencies = [ + "local-waker", + "pin-project-lite", +] + +[[package]] +name = "actix-web" +version = "4.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2e3b15b3dc6c6ed996e4032389e9849d4ab002b1e92fbfe85b5f307d1479b4d" +dependencies = [ + "actix-codec", + "actix-http", + "actix-macros", + "actix-router", + "actix-rt", + "actix-server", + "actix-service", + "actix-utils", + "actix-web-codegen", + "bytes", + "bytestring", + "cfg-if", + "cookie", + "derive_more 2.0.1", + "encoding_rs", + "foldhash", + "futures-core", + "futures-util", + "impl-more", + "itoa", + "language-tags", + "log", + "mime", + "once_cell", + "pin-project-lite", + "regex", + "regex-lite", + "serde", + "serde_json", + "serde_urlencoded", + "smallvec 1.15.1", + "socket2", + "time", + "tracing", + "url", +] + +[[package]] +name = "actix-web-codegen" +version = "4.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f591380e2e68490b5dfaf1dd1aa0ebe78d84ba7067078512b4ea6e4492d622b8" +dependencies = [ + "actix-router", + "proc-macro2", + "quote", + "syn 2.0.102", +] + [[package]] name = "addr2line" version = "0.24.2" @@ -74,6 +257,21 @@ dependencies = [ "equator", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "android-tzdata" version = "0.1.1" @@ -337,6 +535,27 @@ dependencies = [ "generic-array", ] +[[package]] +name = "brotli" +version = "8.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "built" version = "0.7.7" @@ -393,6 +612,15 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +[[package]] +name = "bytestring" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "113b4343b5f6617e7ad401ced8de3cc8b012e73a594347c307b90db3e9271289" +dependencies = [ + "bytes", +] + [[package]] name = "candle-core" version = "0.9.1" @@ -651,6 +879,17 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e" +[[package]] +name = "cookie" +version = "0.16.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e859cd57d0710d9e06c381b550c06e76992472a8c6d527aecd2fc673dcc231fb" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -801,6 +1040,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "deranged" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a41953f86f8a05768a6cda24def994fd2f424b04ec5c719cf89989779f199071" +dependencies = [ + "powerfmt", +] + [[package]] name = "derive_arbitrary" version = "1.4.1" @@ -856,6 +1104,27 @@ dependencies = [ "syn 2.0.102", ] +[[package]] +name = "derive_more" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "093242cf7570c207c83073cf82f79706fe7b8317e98620a47d5be7c3d8497678" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda628edc44c4bb645fbe0f758797143e4e07926f7ebf4e9bdfbd3d2ce621df3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.102", + "unicode-xid", +] + [[package]] name = "derive_utils" version = "0.15.0" @@ -937,7 +1206,7 @@ version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "21507309413f0ba54ec1f2beeda534d116d3db94bffea9c18991bea116f3e522" dependencies = [ - "derive_more", + "derive_more 0.99.20", "hard-xml", "log", "zip", @@ -1189,6 +1458,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "foreign-types" version = "0.3.2" @@ -1626,6 +1901,25 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "half" version = "2.6.0" @@ -1707,7 +2001,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cc03dcb0b0a83ae3f3363ec811014ae669f083e4e499c66602f447c4828737a1" dependencies = [ "dirs", - "http", + "http 1.3.1", "indicatif", "libc", "log", @@ -1744,6 +2038,17 @@ dependencies = [ "syn 2.0.102", ] +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.3.1" @@ -1762,7 +2067,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.3.1", ] [[package]] @@ -1773,7 +2078,7 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", + "http 1.3.1", "http-body", "pin-project-lite", ] @@ -1784,6 +2089,12 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "hyper" version = "1.6.0" @@ -1793,7 +2104,7 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http", + "http 1.3.1", "http-body", "httparse", "itoa", @@ -1809,7 +2120,7 @@ version = "0.27.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ - "http", + "http 1.3.1", "hyper", "hyper-util", "rustls", @@ -1831,7 +2142,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "http", + "http 1.3.1", "http-body", "hyper", "ipnet", @@ -2140,6 +2451,12 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0263a3d970d5c054ed9312c0057b4f3bde9c0b33836d3637361d4a9e6e7a408" +[[package]] +name = "impl-more" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a5a9a0ff0086c7a148acb942baaabeadf9504d10400b5a05645853729b9cd2" + [[package]] name = "indexmap" version = "2.9.0" @@ -2308,6 +2625,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "language-tags" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" + [[package]] name = "lazy_static" version = "1.5.0" @@ -2381,6 +2704,23 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" +[[package]] +name = "local-channel" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6cbc85e69b8df4b8bb8b89ec634e7189099cea8927a276b7384ce5488e53ec8" +dependencies = [ + "futures-core", + "futures-sink", + "local-waker", +] + +[[package]] +name = "local-waker" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d873d7c67ce09b42110d801813efbc9364414e356be9935700d368351657487" + [[package]] name = "lock_api" version = "0.4.13" @@ -2580,6 +2920,12 @@ dependencies = [ "paste", ] +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + [[package]] name = "minimal-lexical" version = "0.2.1" @@ -2603,6 +2949,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78bed444cc8a2160f01cbcf811ef18cac863ad68ae8ca62092e8db51d51c761c" dependencies = [ "libc", + "log", "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.59.0", ] @@ -2800,6 +3147,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + [[package]] name = "num-derive" version = "0.4.2" @@ -3263,6 +3616,12 @@ dependencies = [ "zerovec 0.11.2", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -3896,6 +4255,12 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "regex-lite" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d942b98df5e658f56f20d592c7f868833fe38115e65c33003d8cd224b0155da" + [[package]] name = "regex-syntax" version = "0.8.5" @@ -3922,7 +4287,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "http", + "http 1.3.1", "http-body", "http-body-util", "hyper", @@ -4145,18 +4510,28 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.219" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.219" +version = "1.0.228" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", @@ -4205,6 +4580,26 @@ dependencies = [ "serde", ] +[[package]] +name = "server" +version = "0.6.5" +dependencies = [ + "actix-web", + "embed_anything", + "serde", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sha2" version = "0.10.9" @@ -4222,6 +4617,15 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook-registry" +version = "1.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a4719bff48cee6b39d12c020eeb490953ad2443b7055bd0b21fca26bd8c28b" +dependencies = [ + "libc", +] + [[package]] name = "simd-adler32" version = "0.3.7" @@ -4826,6 +5230,37 @@ dependencies = [ "weezl", ] +[[package]] +name = "time" +version = "0.3.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e7d9e3bb61134e77bde20dd4825b97c010155709965fedf0f49bb138e52a9d" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40868e7c1d2f0b8d73e4a8c7f0ff63af4f6d19be117e90bd73eb1d62cf831c6b" + +[[package]] +name = "time-macros" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30cfb0125f12d9c277f35663a0a33f8c30190f4e4574868a330595412d34ebf3" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.7.6" @@ -4903,7 +5338,9 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.52.0", @@ -5008,7 +5445,7 @@ dependencies = [ "bitflags 2.9.1", "bytes", "futures-util", - "http", + "http 1.3.1", "http-body", "iri-string", "pin-project-lite", @@ -5035,6 +5472,7 @@ version = "0.1.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" dependencies = [ + "log", "pin-project-lite", "tracing-attributes", "tracing-core", @@ -5183,6 +5621,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -5247,7 +5691,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fadf18427d33828c311234884b7ba2afb57143e6e7e69fda7ee883b624661e36" dependencies = [ "base64 0.22.1", - "http", + "http 1.3.1", "httparse", "log", ] @@ -6026,6 +6470,34 @@ dependencies = [ "thiserror 1.0.69", ] +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + [[package]] name = "zune-core" version = "0.4.12" diff --git a/Cargo.toml b/Cargo.toml index c3e8ed65..a3276e0a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,6 +3,7 @@ members = [ "processors", "rust", "python", + "server" ] # Python package needs to be built by maturin. exclude = ["python"] diff --git a/docs/guides/actix_server.md b/docs/guides/actix_server.md new file mode 100644 index 00000000..aefc31ae --- /dev/null +++ b/docs/guides/actix_server.md @@ -0,0 +1,126 @@ +# EmbedAnything OpenAI-Compatible Server + +This server provides an OpenAI-compatible API for generating embeddings using the EmbedAnything library. We choose Actix Server for: + +1. Blazing fast: Consistently ranks among the fastest web frameworks in benchmarks like TechEmpower. +2. Asynchronous by default: Built on Rust’s async/await, enabling efficient I/O-bound workloads. +3. Lightweight & modular: Minimal core with extensible middleware, plugins, and integrations. +4. Type-safe: Strong type guarantees ensure fewer runtime surprises. +5. Production-ready: Stable, mature, and already used in industries like fintech, IoT, and SaaS platforms. + + +For benchmarks between python and rust servers, you check out this blog: https://www.jonvet.com/blog/benchmarking-python-rust-web-servers + +## Features + +- OpenAI-compatible `/v1/embeddings` endpoint +- Support for multiple embedding models (Jina, BERT, etc.) +- Health check endpoint + +## Running the Server + +```bash +cargo run -p server --release +``` + +The server will start on `http://0.0.0.0:8080`. + +## API Usage + +### Create Embeddings + +**Endpoint:** `POST /v1/embeddings` + +**Request:** +```json +{ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["The quick brown fox jumps over the lazy dog"] +} +``` + +**Response:** +```json +{ + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [0.0023064255, -0.009327292, ...] + } + ], + "model": "sentence-transformers/all-MiniLM-L12-v2", + "usage": { + "prompt_tokens": 9, + "total_tokens": 9 + } +} +``` + +### Health Check + +**Endpoint:** `GET /health_check` + +Returns a 200 OK status if the server is running. + +## Supported Models + +The server maps model names to EmbedAnything model architectures: + +- `text-embedding-ada-002` → Jina embeddings +- `text-embedding-3-small` → Jina embeddings +- `text-embedding-3-large` → Jina embeddings +- `text-embedding-ada-001` → BERT embeddings +- Unknown models → Default to Jina embeddings + +## Error Handling + +The API returns OpenAI-compatible error responses: + +```json +{ + "error": { + "message": "Error description", + "type": "error_type", + "code": "error_code" + } +} +``` + +## Example Usage with curl + +```bash +# Create embeddings +curl -X POST http://localhost:8080/v1/embeddings \ + -H "Content-Type: application/json" \ + -d '{ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["Hello world", "How are you?"] + }' + +# Health check +curl http://localhost:8080/health_check +``` + +## Example Usage with Python + +```python +import requests + +# Create embeddings +response = requests.post( + "http://localhost:8080/v1/embeddings", + json={ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["The quick brown fox jumps over the lazy dog"] + } +) + +if response.status_code == 200: + data = response.json() + print(f"Generated {len(data['data'])} embeddings") + print(f"First embedding dimension: {len(data['data'][0]['embedding'])}") +else: + print(f"Error: {response.json()}") +``` diff --git a/rust/examples/audio.rs b/rust/examples/audio.rs index 59da8dfe..add14f5e 100644 --- a/rust/examples/audio.rs +++ b/rust/examples/audio.rs @@ -20,7 +20,6 @@ async fn main() { let bert_model = Arc::new( EmbedderBuilder::new() - .model_architecture("bert") .model_id(Some("sentence-transformers/all-MiniLM-L6-v2")) .revision(None) .token(None) diff --git a/rust/examples/bert.rs b/rust/examples/bert.rs index 2dc9613b..f2e5b6ef 100644 --- a/rust/examples/bert.rs +++ b/rust/examples/bert.rs @@ -9,7 +9,6 @@ use std::{path::PathBuf, time::Instant}; async fn main() { let model = Arc::new( EmbedderBuilder::new() - .model_architecture("jina") .model_id(Some("jinaai/jina-embeddings-v2-small-en")) .revision(None) .token(None) diff --git a/rust/examples/clip.rs b/rust/examples/clip.rs index 97305c17..25b7352f 100644 --- a/rust/examples/clip.rs +++ b/rust/examples/clip.rs @@ -10,7 +10,6 @@ async fn main() { let now = Instant::now(); let model = EmbedderBuilder::new() - .model_architecture("clip") .model_id(Some("google/siglip-base-patch16-224")) .revision(None) .token(None) diff --git a/rust/examples/dino.rs b/rust/examples/dino.rs new file mode 100644 index 00000000..f37927db --- /dev/null +++ b/rust/examples/dino.rs @@ -0,0 +1,86 @@ +use candle_core::{Device, Tensor}; +use embed_anything::{ + embed_file, embed_image_directory, embeddings::embed::{Embedder, EmbedderBuilder} +}; +use std::{path::PathBuf, sync::Arc, time::Instant}; + +#[tokio::main] +async fn main() { + let now = Instant::now(); + + let model = EmbedderBuilder::new() + .model_id(Some("facebook/dinov2-small")) + .revision(None) + .token(None) + .from_pretrained_hf() + .unwrap(); + let model: Arc = Arc::new(model); + // let out = embed_image_directory(PathBuf::from("test_files"), &model, None, None) + // .await + // .unwrap() + // .unwrap(); + + let query_emb_data = embed_file("test_files/clip/cat3.jpg", &model, None, None) + .await + .unwrap().unwrap(); + // println!("query_emb_data: {:?}", query_emb_data[0].embedding); + // let n_vectors = out.len(); + + // let vector = out + // .iter() + // .map(|embed| embed.embedding.clone()) + // .collect::>() + // .into_iter() + // .flat_map(|x| x.to_dense().unwrap()) + // .collect::>(); + + // let out_embeddings = Tensor::from_vec( + // vector, + // (n_vectors, out[0].embedding.to_dense().unwrap().len()), + // &Device::Cpu, + // ) + // .unwrap(); + + // let image_paths = out + // .iter() + // .map(|embed| embed.text.clone().unwrap()) + // .collect::>(); + + // let query_embeddings = Tensor::from_vec( + // query_emb_data + // .iter() + // .map(|embed| embed.embedding.clone()) + // .collect::>() + // .into_iter() + // .flat_map(|x| x.to_dense().unwrap()) + // .collect::>(), + // (1, query_emb_data[0].embedding.to_dense().unwrap().len()), + // &Device::Cpu, + // ) + // .unwrap(); + + // let similarities = out_embeddings + // .matmul(&query_embeddings.transpose(0, 1).unwrap()) + // .unwrap() + // .detach() + // .squeeze(1) + // .unwrap() + // .to_vec1::() + // .unwrap(); + + // let mut indices: Vec = (0..similarities.len()).collect(); + // indices.sort_by(|a, b| similarities[*b].partial_cmp(&similarities[*a]).unwrap()); + + // println!("Descending order of similarity: "); + // for idx in &indices { + // println!("{}", image_paths[*idx]); + // println!("Similarity: {}", similarities[*idx]); + // } + + // println!("-----------"); + + // println!("Most similar image: {}", image_paths[indices[0]]); + + // let elapsed_time = now.elapsed(); + // println!("Elapsed Time: {}", elapsed_time.as_secs_f32()); +} \ No newline at end of file diff --git a/rust/examples/late_chunking.rs b/rust/examples/late_chunking.rs index 03db3b25..3f1bf533 100644 --- a/rust/examples/late_chunking.rs +++ b/rust/examples/late_chunking.rs @@ -6,7 +6,6 @@ use embed_anything::{config::TextEmbedConfig, embeddings::embed::EmbedderBuilder async fn main() { let model = Arc::new( EmbedderBuilder::new() - .model_architecture("jina") .model_id(Some("jinaai/jina-embeddings-v2-small-en")) .revision(None) .path_in_repo(Some("model.onnx")) diff --git a/rust/examples/model2vec.rs b/rust/examples/model2vec.rs index 8587578d..6732fd9c 100644 --- a/rust/examples/model2vec.rs +++ b/rust/examples/model2vec.rs @@ -9,7 +9,6 @@ use std::{path::PathBuf, time::Instant}; async fn main() { let model = Arc::new( EmbedderBuilder::new() - .model_architecture("model2vec") .model_id(Some("minishlab/potion-base-8M")) .revision(None) .token(None) diff --git a/rust/examples/splade.rs b/rust/examples/splade.rs index c22def29..00f14993 100644 --- a/rust/examples/splade.rs +++ b/rust/examples/splade.rs @@ -43,7 +43,6 @@ async fn main() -> anyhow::Result<()> { ), ModelType::Normal => Arc::new( EmbedderBuilder::new() - .model_architecture("sparse-bert") .model_id(Some("prithivida/Splade_PP_en_v1")) .revision(None) .from_pretrained_hf() diff --git a/rust/examples/web_embed.rs b/rust/examples/web_embed.rs index 4f2308a8..88962df0 100644 --- a/rust/examples/web_embed.rs +++ b/rust/examples/web_embed.rs @@ -14,7 +14,6 @@ async fn main() { let embedder = Arc::new( EmbedderBuilder::new() - .model_architecture("bert") .model_id(Some("sentence-transformers/all-MiniLM-L6-v2")) .revision(None) .from_pretrained_hf() diff --git a/rust/src/embeddings/embed.rs b/rust/src/embeddings/embed.rs index 0b701a2d..cd43861d 100644 --- a/rust/src/embeddings/embed.rs +++ b/rust/src/embeddings/embed.rs @@ -1,4 +1,5 @@ use crate::config::{ImageEmbedConfig, TextEmbedConfig}; +use crate::embeddings::local::vision_encoder::VisionEncoderEmbedder; use crate::file_processor::audio::audio_processor::Segment; use crate::Dtype; @@ -16,6 +17,7 @@ use super::local::qwen3::{Qwen3Embed, Qwen3Embedder}; use super::local::text_embedding::ONNXModel; use anyhow::anyhow; use anyhow::Result; +use hf_hub::Repo; use serde::Deserialize; use std::collections::HashMap; use std::future::Future; @@ -139,42 +141,38 @@ impl TextEmbedder { } pub fn from_pretrained_hf( - model: &str, + architecture: &str, model_id: &str, revision: Option<&str>, token: Option<&str>, dtype: Option, ) -> Result { - match model { - "jina" | "Jina" => Ok(Self::Jina(Box::new(JinaEmbedder::new( + match architecture { + "JinaBertForMaskedLM" => Ok(Self::Jina(Box::new(JinaEmbedder::new( model_id, revision, token, )?))), - "Bert" | "bert" => Ok(Self::Bert(Box::new(BertEmbedder::new( + "BertModel" => Ok(Self::Bert(Box::new(BertEmbedder::new( model_id.to_string(), revision.map(|s| s.to_string()), token, )?))), - "sparse-bert" | "SparseBert" | "SPARSE-BERT" => { - Ok(Self::Bert(Box::new(SparseBertEmbedder::new( - model_id.to_string(), - revision.map(|s| s.to_string()), - token, - )?))) - } - "model2vec" | "Model2Vec" | "MODEL2VEC" => Ok(Self::Model2Vec(Box::new( - Model2VecEmbedder::new(model_id, token, None)?, - ))), + "BertForMaskedLM" => Ok(Self::Bert(Box::new(SparseBertEmbedder::new( + model_id.to_string(), + revision.map(|s| s.to_string()), + token, + )?))), + "StaticModel" => Ok(Self::Model2Vec(Box::new(Model2VecEmbedder::new( + model_id, token, None, + )?))), - "modernbert" | "ModernBert" | "MODERNBERT" => { - Ok(Self::ModernBert(Box::new(ModernBertEmbedder::new( - model_id.to_string(), - revision.map(|s| s.to_string()), - token, - dtype, - )?))) - } - "qwen3" | "Qwen3" | "QWEN3" => Ok(Self::Qwen3(Box::new(Qwen3Embedder::new( + "ModernBertForMaskedLM" => Ok(Self::ModernBert(Box::new(ModernBertEmbedder::new( + model_id.to_string(), + revision.map(|s| s.to_string()), + token, + dtype, + )?))), + "Qwen3EmbedForCausalLM" => Ok(Self::Qwen3(Box::new(Qwen3Embedder::new( model_id, revision.map(|s| s.to_string()), token, @@ -275,9 +273,7 @@ impl TextEmbedder { model_id.to_string(), api_key, ))), - "gemini" | "Gemini" => Ok(Self::Gemini(GeminiEmbedder::new( - api_key, - ))), + "gemini" | "Gemini" => Ok(Self::Gemini(GeminiEmbedder::new(api_key))), _ => Err(anyhow::anyhow!("Model not supported")), } } @@ -285,6 +281,7 @@ impl TextEmbedder { pub enum VisionEmbedder { Clip(Box), + VisionEncoder(Box), ColPali(Box), Cohere(CohereEmbedder), } @@ -315,20 +312,26 @@ impl From for TextEmbedder { impl VisionEmbedder { pub fn from_pretrained_hf( - model: &str, + architecture: &str, model_id: &str, revision: Option<&str>, token: Option<&str>, - ) -> Result { - match model { - "clip" | "Clip" | "CLIP" => Ok(Self::Clip(Box::new(ClipEmbedder::new( + ) -> Result { + match architecture { + "CLIPModel" | "SiglipModel" => Ok(Self::Clip(Box::new(ClipEmbedder::new( model_id.to_string(), revision, token, )?))), - "colpali" | "ColPali" | "COLPALI" => Ok(Self::ColPali(Box::new(ColPaliEmbedder::new( + "Dinov2Model" => Ok(Self::VisionEncoder(Box::new(VisionEncoderEmbedder::new( + model_id, + revision, + token, + )?))), + "ColPali" => Ok(Self::ColPali(Box::new(ColPaliEmbedder::new( model_id, revision, )?))), + _ => Err(anyhow::anyhow!("Model not supported")), } } @@ -473,7 +476,6 @@ impl EmbedderBuilder { pub fn from_pretrained_hf(self) -> Result { match self.model_id { Some(model_id) => Embedder::from_pretrained_hf( - &self.model_architecture, &model_id, self.revision.as_deref(), self.token.as_deref(), @@ -538,62 +540,80 @@ impl Embedder { } pub fn from_pretrained_hf( - model_architecture: &str, model_id: &str, revision: Option<&str>, token: Option<&str>, dtype: Option, ) -> Result { - match model_architecture { - "clip" | "Clip" | "CLIP" => Ok(Self::Vision(Box::new( - VisionEmbedder::from_pretrained_hf(model_architecture, model_id, revision, token)?, + let api = hf_hub::api::sync::ApiBuilder::from_env() + .with_token(token.map(|s| s.to_string())) + .build()?; + let api = match revision { + Some(rev) => api.repo(Repo::with_revision( + model_id.to_string(), + hf_hub::RepoType::Model, + rev.to_string(), + )), + None => api.repo(Repo::new(model_id.to_string(), hf_hub::RepoType::Model)), + }; + let config_filename = api.get("config.json")?; + let config = std::fs::read_to_string(config_filename)?; + let config: serde_json::Value = serde_json::from_str(&config)?; + + let architecture = config["architectures"] + .as_array() + .ok_or(anyhow::anyhow!("Architecture not found"))? + .first() + .ok_or(anyhow::anyhow!("Architecture not found"))? + .as_str() + .ok_or(anyhow::anyhow!("Architecture not found"))?; + match architecture { + "CLIPModel" | "SiglipModel" => Ok(Self::Vision(Box::new( + VisionEmbedder::from_pretrained_hf(architecture, model_id, revision, token)?, ))), - "colpali" | "ColPali" | "COLPALI" => Ok(Self::Vision(Box::new( - VisionEmbedder::from_pretrained_hf(model_architecture, model_id, revision, token)?, + "ColPali" => Ok(Self::Vision(Box::new( + VisionEmbedder::from_pretrained_hf(architecture, model_id, revision, token)?, ))), - "bert" | "Bert" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, + "Dinov2Model" => Ok(Self::Vision(Box::new( + VisionEmbedder::from_pretrained_hf(architecture, model_id, revision, token)?, + ))), + "BertModel" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, model_id, revision, token, dtype, )?)), - "jina" | "Jina" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, + "JinaBertForMaskedLM" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, model_id, revision, token, dtype, )?)), - "model2vec" | "Model2Vec" | "MODEL2VEC" => { - Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, - model_id, - revision, - token, - dtype, - )?)) - } - "sparse-bert" | "SparseBert" | "SPARSE-BERT" => { - Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, - model_id, - revision, - token, - dtype, - )?)) - } - "modernbert" | "ModernBert" | "MODERNBERT" => { - Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, - model_id, - revision, - token, - dtype, - )?)) - } - "qwen3" | "Qwen3" | "QWEN3" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( - model_architecture, + "StaticModel" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, + model_id, + revision, + token, + dtype, + )?)), + "BertForMaskedLM" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, + model_id, + revision, + token, + dtype, + )?)), + "ModernBertForMaskedLM" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, + model_id, + revision, + token, + dtype, + )?)), + "Qwen3EmbedForCausalLM" => Ok(Self::Text(TextEmbedder::from_pretrained_hf( + architecture, model_id, revision, token, @@ -801,6 +821,7 @@ impl TextEmbed for VisionEmbedder { ) -> Result, anyhow::Error> { match self { Self::Clip(embedder) => embedder.embed(text_batch, batch_size), + Self::VisionEncoder(_) => Err(anyhow::anyhow!("Model not supported for text embedding")), Self::ColPali(embedder) => embedder.embed(text_batch, batch_size), Self::Cohere(embedder) => embedder.embed(text_batch).await, } @@ -837,6 +858,7 @@ impl EmbedImage for VisionEmbedder { embedder.embed_image(PathBuf::from(image_path.as_ref()), metadata) } Self::Cohere(embedder) => embedder.embed_image(image_path, metadata).await, + Self::VisionEncoder(embedder) => embedder.embed_image(image_path, metadata).await, } } @@ -854,6 +876,7 @@ impl EmbedImage for VisionEmbedder { .collect::>(), ), Self::Cohere(embedder) => embedder.embed_image_batch(image_paths, batch_size).await, + Self::VisionEncoder(embedder) => embedder.embed_image_batch(image_paths, batch_size).await, } } diff --git a/rust/src/embeddings/local/clip.rs b/rust/src/embeddings/local/clip.rs index fa56e900..b1e5357c 100644 --- a/rust/src/embeddings/local/clip.rs +++ b/rust/src/embeddings/local/clip.rs @@ -17,8 +17,8 @@ use crate::{ }, }; use candle_core::{DType, Device, Tensor}; - -use candle_nn::VarBuilder; +use candle_nn::{Module, VarBuilder}; +use candle_transformers::models::dinov2::DinoVisionTransformer; use tokenizers::{PaddingParams, Tokenizer}; use crate::embeddings::embed::{EmbedData, EmbedImage}; @@ -26,6 +26,7 @@ use crate::embeddings::embed::{EmbedData, EmbedImage}; pub enum VisionModel { Clip(clip::ClipModel), Siglip(siglip::Model), + Dino(DinoVisionTransformer), } impl VisionModel { @@ -33,6 +34,7 @@ impl VisionModel { match self { VisionModel::Clip(model) => model.get_image_features(image), VisionModel::Siglip(model) => model.get_image_features(image), + VisionModel::Dino(model) => model.forward(image), } } @@ -40,6 +42,17 @@ impl VisionModel { match self { VisionModel::Clip(model) => model.get_text_features(text), VisionModel::Siglip(model) => model.get_text_features(text), + VisionModel::Dino(_) => Err(candle_core::Error::msg( + "Dino model does not support text features", + )), + } + } + + pub fn supports_text(&self) -> bool { + match self { + VisionModel::Clip(_) => true, + VisionModel::Siglip(_) => true, + VisionModel::Dino(_) => false, } } } @@ -100,6 +113,7 @@ impl ClipEmbedder { let config_str = std::fs::read_to_string(config_filename)?; let config_json: serde_json::Value = serde_json::from_str(&config_str)?; + let mut tokenizer = Self::get_tokenizer(None, model_id, revision)?; let pp = PaddingParams { strategy: tokenizers::PaddingStrategy::BatchLongest, @@ -130,6 +144,35 @@ impl ClipEmbedder { config.text_config.pad_token_id as u32, ) } + Some("Dinov2Model") => { + let (depth, embed_dim, num_heads) = ( + config_json + .get("num_hidden_layers") + .expect("num_hidden_layers not found") + .as_u64() + .expect("num_hidden_layers not a number") as usize, + config_json + .get("hidden_size") + .expect("hidden_size not found") + .as_u64() + .expect("hidden_size not a number") as usize, + config_json + .get("num_attention_heads") + .expect("num_attention_heads not found") + .as_u64() + .expect("num_attention_heads not a number") as usize, + ); + ( + VisionModel::Dino(DinoVisionTransformer::new( + vb, + depth as usize, + embed_dim as usize, + num_heads as usize, + )?), + 0, + 0, + ) + } _ => return Err(anyhow::Error::msg("Unsupported model architecture")), } } else { @@ -260,6 +303,10 @@ impl ClipEmbedder { ) -> Result, anyhow::Error> { let mut encodings = Vec::new(); + if !self.model.supports_text() { + return Err(anyhow::Error::msg("This model does not support text features")); + } + let batch_size = batch_size.unwrap_or(32); for mini_text_batch in text_batch.chunks(batch_size) { diff --git a/rust/src/embeddings/local/mod.rs b/rust/src/embeddings/local/mod.rs index d071412b..d9eefd04 100644 --- a/rust/src/embeddings/local/mod.rs +++ b/rust/src/embeddings/local/mod.rs @@ -21,4 +21,5 @@ pub mod ort_bert; pub mod ort_jina; pub mod pooling; pub mod qwen3; +pub mod vision_encoder; pub mod text_embedding; diff --git a/rust/src/embeddings/local/ort_bert.rs b/rust/src/embeddings/local/ort_bert.rs index 86a8eb6f..e6b0845a 100644 --- a/rust/src/embeddings/local/ort_bert.rs +++ b/rust/src/embeddings/local/ort_bert.rs @@ -85,15 +85,15 @@ impl OrtBertEmbedder { base_path.push('/'); } let model_path = match dtype { - Some(Dtype::Q4F16) => format!("{base_path}/model_q4f16.onnx"), - Some(Dtype::F16) => format!("{base_path}/model_fp16.onnx"), - Some(Dtype::INT8) => format!("{base_path}/model_int8.onnx"), - Some(Dtype::Q4) => format!("{base_path}/model_q4.onnx"), - Some(Dtype::UINT8) => format!("{base_path}/model_uint8.onnx"), - Some(Dtype::BNB4) => format!("{base_path}/model_bnb4.onnx"), - Some(Dtype::F32) => format!("{base_path}/model.onnx"), - Some(Dtype::QUANTIZED) => format!("{base_path}/model_quantized.onnx"), - Some(Dtype::BF16) => format!("{base_path}/model_bf16.onnx"), + Some(Dtype::Q4F16) => format!("{base_path}model_q4f16.onnx"), + Some(Dtype::F16) => format!("{base_path}model_fp16.onnx"), + Some(Dtype::INT8) => format!("{base_path}model_int8.onnx"), + Some(Dtype::Q4) => format!("{base_path}model_q4.onnx"), + Some(Dtype::UINT8) => format!("{base_path}model_uint8.onnx"), + Some(Dtype::BNB4) => format!("{base_path}model_bnb4.onnx"), + Some(Dtype::F32) => format!("{base_path}model.onnx"), + Some(Dtype::QUANTIZED) => format!("{base_path}model_quantized.onnx"), + Some(Dtype::BF16) => format!("{base_path}model_bf16.onnx"), None => path.to_string(), }; let weights = api.get(model_path.as_str()); diff --git a/rust/src/embeddings/local/qwen3.rs b/rust/src/embeddings/local/qwen3.rs index 668d14d8..3f9d7345 100644 --- a/rust/src/embeddings/local/qwen3.rs +++ b/rust/src/embeddings/local/qwen3.rs @@ -121,15 +121,12 @@ impl Qwen3Embed for Qwen3Embedder { for mini_text_batch in text_batch.chunks(batch_size) { let (token_ids, attention_mask) = tokenize_batch(&self.tokenizer, mini_text_batch, &self.device)?; - let embeddings: Tensor = self - .model - .write() - .unwrap() - .forward(&token_ids, &attention_mask, 0) - .unwrap() - .to_dtype(DType::F32)?; - - self.model.write().unwrap().clear_kv_cache(); + let embeddings: Tensor = { + let mut model = self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?; + model.forward(&token_ids, &attention_mask, 0)?.to_dtype(DType::F32)? + }; + + self.model.write().map_err(|e| anyhow::anyhow!("Lock poisoned: {}", e))?.clear_kv_cache(); let attention_mask = PooledOutputType::from(attention_mask); let attention_mask = Some(&attention_mask); let model_output = ModelOutput::Tensor(embeddings.clone()); diff --git a/rust/src/embeddings/local/vision_encoder.rs b/rust/src/embeddings/local/vision_encoder.rs new file mode 100644 index 00000000..131feff1 --- /dev/null +++ b/rust/src/embeddings/local/vision_encoder.rs @@ -0,0 +1,302 @@ +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use std::{collections::HashMap, fs}; + +use anyhow::Error as E; + +use crate::{ + embeddings::{embed::EmbeddingResult, select_device}, + models::clip::div_l2_norm, +}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::{Module, VarBuilder}; + +use crate::models::dinov2::DinoVisionTransformer; +use crate::embeddings::embed::{EmbedData, EmbedImage}; + +pub enum VisionEncoderModel { + Dino(DinoVisionTransformer), +} + +impl VisionEncoderModel { + pub fn get_image_features(&self, image: &Tensor) -> Result { + match self { + VisionEncoderModel::Dino(model) => model.forward(image), + } + } +} + +pub struct VisionEncoderEmbedder { + pub model: VisionEncoderModel, + pub device: Device, + pub image_size: usize, +} +impl Default for VisionEncoderEmbedder { + fn default() -> Self { + Self::new("facebook/dinov2-small", None, None).unwrap() + } +} + +impl VisionEncoderEmbedder { + pub fn new(model_id: &str, revision: Option<&str>, token: Option<&str>) -> Result { + let api = hf_hub::api::sync::ApiBuilder::from_env() + .with_token(token.map(|s| s.to_string())) + .build()?; + + let api = match revision { + Some(rev) => api.repo(hf_hub::Repo::with_revision( + model_id.to_string(), + hf_hub::RepoType::Model, + rev.to_string(), + )), + None => api.repo(hf_hub::Repo::new( + model_id.to_string(), + hf_hub::RepoType::Model, + )), + }; + + let device = select_device(); + + let vb = match api.get("model.safetensors") { + Ok(safetensors) => unsafe { + VarBuilder::from_mmaped_safetensors(&[safetensors], DType::F32, &device)? + }, + Err(_) => match api.get("pytorch_model.bin") { + Ok(pytorch_model) => VarBuilder::from_pth(pytorch_model, DType::F32, &device)?, + Err(e) => { + return Err(anyhow::Error::msg(format!( + "Model weights not found. The weights should either be a `model.safetensors` or `pytorch_model.bin` file. Error: {}", + e + ))); + } + }, + }; + let config_filename = api.get("config.json")?; + let config_str = std::fs::read_to_string(config_filename)?; + let config_json: serde_json::Value = serde_json::from_str(&config_str)?; + + let image_size = config_json + .get("image_size") + .expect("image_size not found") + .as_u64() + .expect("image_size not a number") as usize; + + let model = if let Some(architectures) = config_json.get("architectures") { + if let Some(arch) = architectures.get(0) { + match arch.as_str() { + Some("Dinov2Model") => { + let (depth, embed_dim, num_heads) = ( + config_json + .get("num_hidden_layers") + .expect("num_hidden_layers not found") + .as_u64() + .expect("num_hidden_layers not a number") + as usize, + config_json + .get("hidden_size") + .expect("hidden_size not found") + .as_u64() + .expect("hidden_size not a number") + as usize, + config_json + .get("num_attention_heads") + .expect("num_attention_heads not found") + .as_u64() + .expect("num_attention_heads not a number") + as usize, + ); + + VisionEncoderModel::Dino(DinoVisionTransformer::new( + vb, + depth as usize, + embed_dim as usize, + num_heads as usize, + )?) + } + _ => return Err(anyhow::Error::msg("Unsupported model architecture")), + } + } else { + return Err(anyhow::Error::msg("No architecture specified in config")); + } + } else { + VisionEncoderModel::Dino(DinoVisionTransformer::new(vb, 12, 384, 6)?) + }; + Ok(VisionEncoderEmbedder { + model, + device, + image_size, + }) + } + + fn load_image>( + &self, + path: T, + image_size: usize, + ) -> anyhow::Result { + let img = image::ImageReader::open(path)?.decode()?; + let (height, width) = (224, 224); + let img = img.resize_exact( + width as u32, + height as u32, + image::imageops::FilterType::CatmullRom, + ); + + let img = img.to_rgb8(); + + let img = img.into_raw(); + let img = Tensor::from_vec(img, (height, width, 3), &self.device)? + .permute((2, 0, 1))? + .to_dtype(DType::F32)? + .affine(0.00392156862745098, 0.)?; // Rescale factor: 1/255 + println!("img: {}", img); + // ImageNet normalization: (image - mean) / std + let mean = Tensor::new(&[0.485f32, 0.456f32, 0.406f32], &self.device)? + .unsqueeze(1)? + .unsqueeze(2)?; // Shape: [3, 1, 1] + let std = Tensor::new(&[0.229f32, 0.224f32, 0.225f32], &self.device)? + .unsqueeze(1)? + .unsqueeze(2)?; // Shape: [3, 1, 1] + + let img = img.broadcast_sub(&mean)?.broadcast_div(&std)?; + // .unsqueeze(0)?; + Ok(img) + } + + fn load_images>( + &self, + paths: &[T], + image_size: usize, + ) -> anyhow::Result { + let mut images = vec![]; + + for path in paths { + let tensor = self.load_image(path, image_size)?; + images.push(tensor); + } + + let images = Tensor::stack(&images, 0)?; + + Ok(images) + } +} + +impl EmbedImage for VisionEncoderEmbedder { + async fn embed_image_batch>( + &self, + image_paths: &[T], + batch_size: Option, + ) -> anyhow::Result> { + let mut encodings = Vec::new(); + for image_batch in image_paths.chunks(batch_size.unwrap_or(32)) { + let images = self.load_images(image_batch, self.image_size).unwrap(); + let batch_encodings = self.model.get_image_features(&images)?; + let normalized_encodings = div_l2_norm(&batch_encodings)?.to_vec2::()?; + encodings.extend(normalized_encodings); + } + + let embeddings = encodings + .iter() + .zip(image_paths) + .map(|(data, path)| { + let mut metadata = HashMap::new(); + metadata.insert( + "file_name".to_string(), + fs::canonicalize(path) + .unwrap() + .to_str() + .unwrap() + .to_string(), + ); + + EmbedData::new( + EmbeddingResult::DenseVector(data.to_vec()), + Some(path.as_ref().to_str().unwrap().to_string()), + Some(metadata), + ) + }) + .collect::>(); + Ok(embeddings) + } + + async fn embed_image>( + &self, + image_path: T, + metadata: Option>, + ) -> anyhow::Result { + let image = self + .load_image(&image_path, self.image_size) + .unwrap() + .unsqueeze(0) + .unwrap(); + let encoding = &self.model.get_image_features(&image)?; + let normalized_encoding = &div_l2_norm(encoding)?.to_vec2::()?[0]; + Ok(EmbedData::new( + EmbeddingResult::DenseVector(normalized_encoding.to_vec()), + None, + metadata.clone(), + )) + } + + async fn embed_pdf>( + &self, + _pdf_path: T, + _batch_size: Option, + ) -> anyhow::Result> { + Err(anyhow::anyhow!( + "PDF embedding not supported for Clip model" + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Tests the load_image method. + #[test] + fn test_load_image() { + let vision_encoder_embedder = VisionEncoderEmbedder::default(); + let image = vision_encoder_embedder + .load_image("../test_files/clip/cat3.jpg", 224) + .unwrap(); + println!("image: {}", image); + assert_eq!(image.shape().clone().into_dims(), &[3, 224, 224]); + } + + // Tests the load_images method. + #[test] + fn test_load_images() { + let vision_encoder_embedder = VisionEncoderEmbedder::default(); + let images = vision_encoder_embedder + .load_images( + &[ + "../test_files/clip/cat1.jpg", + "../test_files/clip/cat2.jpeg", + ], + 224, + ) + .unwrap(); + assert_eq!(images.shape().clone().into_dims(), &[2, 3, 224, 224]); + } + + // Tests the embed_image_batch method. + #[tokio::test] + async fn test_embed_image_batch() { + let vision_encoder_embedder = VisionEncoderEmbedder::default(); + let embeddings = vision_encoder_embedder + .embed_image_batch( + &[ + "../test_files/clip/cat1.jpg", + "../test_files/clip/cat2.jpeg", + ], + Some(2), + ) + .await + .unwrap(); + assert_eq!(embeddings.len(), 2); + } +} diff --git a/rust/src/models/dinov2.rs b/rust/src/models/dinov2.rs new file mode 100644 index 00000000..c592b48a --- /dev/null +++ b/rust/src/models/dinov2.rs @@ -0,0 +1,452 @@ +//! Implementation of the DINOv2 models from Meta Research. +//! +//! This module implements the DINOv2 vision transformer model from Meta AI Research. +//! DINOv2 is a self-supervised learning model that can learn visual features +//! without using any labeled data. See: ["DINOv2: Learning Robust Visual Features without Supervision"](https://github.com/facebookresearch/dinov2) +//! +//! ## Running an example with color map and CUDA +//! +//! ```bash +//! cargo run \ +//! --features cuda,depth_anything_v2 \ +//! --package candle-examples \ +//! --example depth_anything_v2 \ +//! -- --color-map \ +//! --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! ``` +//! +//! ## Running as an ImageNet classifier +//! +//! The model returns the probability for the image to belong to each of the 1000 ImageNet categories. +//! +//!
+//! +//!
+//! +//! ```bash +//! cargo run \ +//! --example dinov2 \ +//! --release \ +//! -- --image candle-examples/examples/yolo-v8/assets/bike.jpg +//! +//! > mountain bike, all-terrain bike, off-roader: 43.67% +//! > bicycle-built-for-two, tandem bicycle, tandem: 33.20% +//! > crash helmet : 13.23% +//! > unicycle, monocycle : 2.44% +//! > maillot : 2.42% +//! ``` +//! + +use candle_core::{IndexOp, Result, Tensor, D}; +use candle_nn::{layer_norm, LayerNorm, Linear, Module, VarBuilder}; + +const IMG_SIZE: usize = 518; +const PATCH_SIZE: usize = 14; + +fn linear(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result { + if bias { + candle_nn::linear(in_dim, out_dim, vb) + } else { + candle_nn::linear_no_bias(in_dim, out_dim, vb) + } +} + +#[derive(Debug)] +struct SelfAttention { + query: Linear, + key: Linear, + value: Linear, + num_heads: usize, + head_dim: usize, + scale: f64, +} + +impl SelfAttention { + fn new(vb: VarBuilder, dim: usize, num_heads: usize, qkv_bias: bool) -> Result { + let query = linear(vb.pp("query"), dim, dim, qkv_bias)?; + let key = linear(vb.pp("key"), dim, dim, qkv_bias)?; + let value = linear(vb.pp("value"), dim, dim, qkv_bias)?; + let scale = 1. / ((dim / num_heads) as f64).sqrt(); + let head_dim = dim / num_heads; + Ok(Self { + query, + key, + value, + num_heads, + scale, + head_dim, + }) + } +} + +impl Module for SelfAttention { + fn forward(&self, xs: &Tensor) -> Result { + let (batch_size, q_len, _) = xs.dims3()?; + let query_states = xs.apply(&self.query)?; + let key_states = xs.apply(&self.key)?; + let value_states = xs.apply(&self.value)?; + + let shape = (batch_size, q_len, self.num_heads, self.head_dim); + let query_states = query_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let key_states = key_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + let value_states = value_states.reshape(shape)?.transpose(1, 2)?.contiguous()?; + + let attn_weights = (query_states.matmul(&key_states.t()?)? * self.scale)?; + + // The original implementation upcasts to f32 but candle_nn::ops::softmax should handle this properly. + let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?; + let attn_outputs = attn_weights + .matmul(&value_states)? + .transpose(1, 2)? + .reshape((batch_size, q_len, ()))?; + Ok(attn_outputs) + } +} + +#[derive(Debug)] +struct SelfOutput { + dense: Linear, +} + +impl SelfOutput { + fn new(vb: VarBuilder, hidden_size: usize) -> Result { + let dense = linear(vb.pp("dense"), hidden_size, hidden_size, true)?; + Ok(Self { dense }) + } +} + +impl Module for SelfOutput { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.dense.forward(xs)?; + Ok(xs) + } +} + +#[derive(Debug)] +struct Attention { + self_attention: SelfAttention, + self_output: SelfOutput, +} + +impl Attention { + fn new(vb: VarBuilder, hidden_size: usize, num_heads: usize) -> Result { + let self_attention = SelfAttention::new(vb.pp("attention"), hidden_size, num_heads, true)?; + let self_output = SelfOutput::new(vb.pp("output"), hidden_size)?; + Ok(Self { + self_attention, + self_output, + }) + } +} + +impl Module for Attention { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.self_attention.forward(xs)?; + let xs = self.self_output.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug)] +struct LayerScale { + gamma: Tensor, +} + +impl LayerScale { + fn new(vb: VarBuilder, dim: usize) -> Result { + let gamma = vb.get(dim, "lambda1")?; + Ok(Self { gamma }) + } +} + +impl Module for LayerScale { + fn forward(&self, xs: &Tensor) -> Result { + xs.broadcast_mul(&self.gamma) + } +} + +#[derive(Debug)] +struct Mlp { + fc1: Linear, + fc2: Linear, +} + +impl Mlp { + fn new(vb: VarBuilder, in_features: usize, hidden_features: usize, bias: bool) -> Result { + let out_features = in_features; + let fc1 = linear(vb.pp("fc1"), in_features, hidden_features, bias)?; + let fc2 = linear(vb.pp("fc2"), hidden_features, out_features, bias)?; + Ok(Self { fc1, fc2 }) + } +} + +impl Module for Mlp { + fn forward(&self, xs: &Tensor) -> Result { + let xs = self.fc1.forward(xs)?.gelu()?; + let xs = self.fc2.forward(&xs)?; + Ok(xs) + } +} + +#[derive(Debug)] +struct Block { + norm1: LayerNorm, + attn: Attention, + ls1: LayerScale, + norm2: LayerNorm, + mlp: Mlp, + ls2: LayerScale, +} + +impl Block { + fn new(vb: VarBuilder, dim: usize, num_heads: usize) -> Result { + let norm1 = layer_norm(dim, 1e-5, vb.pp("norm1"))?; + let attn = Attention::new(vb.pp("attention"), dim, num_heads)?; + let ls1 = LayerScale::new(vb.pp("layer_scale1"), dim)?; + let norm2 = layer_norm(dim, 1e-5, vb.pp("norm2"))?; + let mlp = Mlp::new(vb.pp("mlp"), dim, dim * 4, true)?; + let ls2 = LayerScale::new(vb.pp("layer_scale2"), dim)?; + Ok(Self { + norm1, + attn, + ls1, + norm2, + mlp, + ls2, + }) + } +} + +impl Module for Block { + fn forward(&self, xs: &Tensor) -> Result { + let residual = xs; + let xs = self + .ls1 + .forward(&self.attn.forward(&self.norm1.forward(xs)?)?)?; + let xs = (xs + residual)?; + let residual = &xs; + let xs = self + .ls2 + .forward(&self.mlp.forward(&self.norm2.forward(&xs)?)?)?; + xs + residual + } +} + +#[derive(Debug)] +struct PatchEmbed { + proj: candle_nn::Conv2d, + patch_size: (usize, usize), + num_patches: usize, +} + +impl PatchEmbed { + fn new( + vb: VarBuilder, + img_size: usize, + patch_size: usize, + in_chans: usize, + embed_dim: usize, + ) -> Result { + let config = candle_nn::Conv2dConfig { + stride: patch_size, + ..Default::default() + }; + let proj = candle_nn::conv2d(in_chans, embed_dim, patch_size, config, vb.pp("projection"))?; + let num_patches = (img_size / patch_size) * (img_size / patch_size); + Ok(Self { + proj, + patch_size: (patch_size, patch_size), + num_patches, + }) + } +} + +impl Module for PatchEmbed { + fn forward(&self, xs: &Tensor) -> Result { + let (_b, _c, h, w) = xs.dims4()?; + let (patch_h, patch_w) = self.patch_size; + if (h % patch_h) != 0 { + candle_core::bail!("image height {h} is not a multiple of patch height {patch_h}") + } + if (w % patch_w) != 0 { + candle_core::bail!("image width {w} is not a multiple of patch width {patch_w}") + } + let xs = self.proj.forward(xs)?; + let (b, c, h, w) = xs.dims4()?; + // flatten embeddings. + xs.reshape((b, c, h * w))?.transpose(1, 2) + } +} + +#[derive(Debug)] +pub struct DinoVisionTransformer { + patch_embed: PatchEmbed, + cls_token: Tensor, + pos_embed: Tensor, + blocks: Vec, + norm: LayerNorm, +} + +impl DinoVisionTransformer { + pub fn new(vb: VarBuilder, depth: usize, embed_dim: usize, num_heads: usize) -> Result { + let patch_embed = PatchEmbed::new( + vb.pp("embeddings.patch_embeddings"), + IMG_SIZE, + PATCH_SIZE, + 3, + embed_dim, + )?; + let cls_token = vb.get((1, 1, embed_dim), "embeddings.cls_token")?; + let num_tokens = 1; + let pos_embed = vb.get( + (1, patch_embed.num_patches + num_tokens, embed_dim), + "embeddings.position_embeddings", + )?; + let norm = layer_norm(embed_dim, 1e-5, vb.pp("layernorm"))?; + let vb_b = vb.pp("encoder.layer"); + let blocks = (0..depth) + .map(|i| Block::new(vb_b.pp(i.to_string()), embed_dim, num_heads)) + .collect::>>()?; + Ok(Self { + patch_embed, + cls_token, + pos_embed, + blocks, + norm, + }) + } + + fn interpolate_pos_encoding(&self, xs: &Tensor, w: usize, h: usize) -> Result { + let npatch = xs.dim(1)? - 1; + let n = self.pos_embed.dim(1)? - 1; + let sqrt_n = (n as f64).sqrt(); + if npatch == n && w == h { + return Ok(xs.clone()); + } + let class_pos_embed = self.pos_embed.i((.., ..1))?; + let patch_pos_embed = self.pos_embed.i((.., 1..))?; + let dim = xs.dim(D::Minus1)?; + let (w0, h0) = ((w / PATCH_SIZE) as f64 + 0.1, (h / PATCH_SIZE) as f64 + 0.1); + let patch_pos_embed = patch_pos_embed + .reshape((1, sqrt_n as usize, sqrt_n as usize, dim))? + .transpose(2, 3)? + .transpose(1, 2)?; + // This uses bicubic interpolation in the original implementation. + let patch_pos_embed = patch_pos_embed.upsample_nearest2d(h0 as usize, w0 as usize)?; + let el_count = patch_pos_embed.shape().elem_count(); + let patch_pos_embed = + patch_pos_embed + .transpose(1, 2)? + .transpose(2, 3)? + .reshape((1, el_count / dim, dim))?; + Tensor::cat(&[&class_pos_embed, &patch_pos_embed], 1) + } + + fn prepare_tokens_with_mask(&self, xs: &Tensor) -> Result { + let (_b, _nc, w, h) = xs.dims4()?; + let xs = self.patch_embed.forward(xs)?; + + // Repeat (tile) cls_token along the batch dimension to match xs batch size before concat: + let bsz = xs.dim(0)?; + let cls_token = self.cls_token.repeat((bsz, 1, 1))?; + let xs = Tensor::cat(&[&cls_token, &xs], 1)?; + Tensor::broadcast_add(&xs, &self.interpolate_pos_encoding(&xs, w, h)?) + } + + fn get_intermediate_layers_not_chunked( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + ) -> Result> { + let mut xs = self.prepare_tokens_with_mask(xs)?; + let mut output = Vec::new(); + for (i, blk) in self.blocks.iter().enumerate() { + xs = blk.forward(&xs)?; + if blocks_to_take.contains(&i) { + output.push(xs.clone()); + } + } + if output.len() != blocks_to_take.len() { + candle_core::bail!( + "only {} / {} blocks found", + output.len(), + blocks_to_take.len() + ); + } + Ok(output) + } + + pub fn get_intermediate_layers( + &self, + xs: &Tensor, + blocks_to_take: &[usize], + reshape: bool, + return_class_token: bool, + norm: bool, + ) -> Result { + let outputs = self.get_intermediate_layers_not_chunked(xs, blocks_to_take)?; + let outputs = if norm { + outputs + .iter() + .map(|out| self.norm.forward(out)) + .collect::>>()? + } else { + outputs + }; + let class_tokens = outputs + .iter() + .map(|out| out.i((.., 0))) + .collect::>>()?; + let outputs = outputs + .iter() + .map(|out| out.i((.., 1..))) + .collect::>>()?; + + let outputs = if reshape { + let (b, _c, w, h) = xs.dims4()?; + let patch_size = self.patch_embed.patch_size.0; + let num_channels = outputs[0].elem_count() / (b * (w / patch_size) * (h / patch_size)); + outputs + .iter() + .map(|out| { + out.reshape((b, w / patch_size, h / patch_size, num_channels))? + .transpose(2, 3)? + .transpose(1, 2) + }) + .collect::>>()? + } else { + outputs + }; + + let outputs = if return_class_token { + outputs + .iter() + .zip(class_tokens.iter()) + .map(|(out, class_token)| Tensor::cat(&[out, class_token], D::Minus1)) + .collect::>>()? + } else { + outputs + }; + + Tensor::stack(&outputs[..], 0) + } +} + +impl Module for DinoVisionTransformer { + fn forward(&self, xs: &Tensor) -> Result { + let mut xs = self.prepare_tokens_with_mask(xs)?; + for blk in self.blocks.iter() { + xs = blk.forward(&xs)? + } + let xs = self.norm.forward(&xs)?; + let xs_norm_clstoken = xs.i((.., 0, ..))?; + // let xs_norm_patchtokens = xs.i((.., 1.., ..))?.mean(1)?; + + // let xs = Tensor::cat(&[xs_norm_clstoken, xs_norm_patchtokens], D::Minus2)?; + Ok(xs_norm_clstoken) + } +} + +pub fn vit_small(vb: VarBuilder) -> Result { + DinoVisionTransformer::new(vb, 12, 384, 6) +} diff --git a/rust/src/models/mod.rs b/rust/src/models/mod.rs index 5ba30277..695a15c0 100644 --- a/rust/src/models/mod.rs +++ b/rust/src/models/mod.rs @@ -19,6 +19,7 @@ pub mod bert; pub mod clip; pub mod colpali; +pub mod dinov2; pub mod gemma; pub mod idefics3; pub mod jina_bert; diff --git a/server/Cargo.toml b/server/Cargo.toml new file mode 100644 index 00000000..a700b766 --- /dev/null +++ b/server/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "server" +version.workspace = true +edition.workspace = true +license.workspace = true +description.workspace = true +repository.workspace = true +authors.workspace = true + +[dependencies] +actix-web = "4" +serde = { version = "1.0", features = ["derive"] } +embed_anything = {path = "../rust", features = ["ort"]} + +[[bin]] +name = "server" +path = "src/main.rs" \ No newline at end of file diff --git a/server/README.md b/server/README.md new file mode 100644 index 00000000..5cd8d444 --- /dev/null +++ b/server/README.md @@ -0,0 +1,118 @@ +# EmbedAnything OpenAI-Compatible Server + +This server provides an OpenAI-compatible API for generating embeddings using the EmbedAnything library. + +## Features + +- OpenAI-compatible `/v1/embeddings` endpoint +- Support for multiple embedding models (Jina, BERT, etc.) +- Health check endpoint + +## Running the Server + +```bash +cd server +cargo run +``` + +The server will start on `http://0.0.0.0:8080`. + +## API Usage + +### Create Embeddings + +**Endpoint:** `POST /v1/embeddings` + +**Request:** +```json +{ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["The quick brown fox jumps over the lazy dog"] +} +``` + +**Response:** +```json +{ + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [0.0023064255, -0.009327292, ...] + } + ], + "model": "sentence-transformers/all-MiniLM-L12-v2", + "usage": { + "prompt_tokens": 9, + "total_tokens": 9 + } +} +``` + +### Health Check + +**Endpoint:** `GET /health_check` + +Returns a 200 OK status if the server is running. + +## Supported Models + +The server maps model names to EmbedAnything model architectures: + +- `text-embedding-ada-002` → Jina embeddings +- `text-embedding-3-small` → Jina embeddings +- `text-embedding-3-large` → Jina embeddings +- `text-embedding-ada-001` → BERT embeddings +- Unknown models → Default to Jina embeddings + +## Error Handling + +The API returns OpenAI-compatible error responses: + +```json +{ + "error": { + "message": "Error description", + "type": "error_type", + "code": "error_code" + } +} +``` + +## Example Usage with curl + +```bash +# Create embeddings +curl -X POST http://localhost:8080/v1/embeddings \ + -H "Content-Type: application/json" \ + -d '{ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["Hello world", "How are you?"] + }' + +# Health check +curl http://localhost:8080/health_check +``` + +## Example Usage with Python + +```python +import requests + +# Create embeddings +response = requests.post( + "http://localhost:8080/v1/embeddings", + json={ + "model": "sentence-transformers/all-MiniLM-L12-v2", + "input": ["The quick brown fox jumps over the lazy dog"] + } +) + +if response.status_code == 200: + data = response.json() + print(f"Generated {len(data['data'])} embeddings") + print(f"First embedding dimension: {len(data['data'][0]['embedding'])}") +else: + print(f"Error: {response.json()}") +``` diff --git a/server/src/lib.rs b/server/src/lib.rs new file mode 100644 index 00000000..b63ae830 --- /dev/null +++ b/server/src/lib.rs @@ -0,0 +1,150 @@ +use std::net::TcpListener; + +use actix_web::dev::Server; +use actix_web::{get, post, web, App, HttpResponse, HttpServer}; +use serde::{Deserialize, Serialize}; +use embed_anything::config::TextEmbedConfig; +use embed_anything::{embed_query, embeddings::embed::EmbedderBuilder}; + +// OpenAI-compatible request structure +#[derive(Deserialize)] +struct OpenAIEmbedRequest { + model: String, + input: Vec, +} + +// OpenAI-compatible response structures +#[derive(Serialize)] +struct OpenAIEmbedResponse { + object: String, + data: Vec, + model: String, + usage: Usage, +} + +#[derive(Serialize)] +struct EmbeddingData { + object: String, + index: usize, + embedding: Vec, +} + +#[derive(Serialize)] +struct Usage { + prompt_tokens: usize, + total_tokens: usize, +} + +#[derive(Serialize)] +struct ErrorResponse { + error: ErrorDetail, +} + +#[derive(Serialize)] +struct ErrorDetail { + message: String, + #[serde(rename = "type")] + error_type: String, + code: Option, +} + +#[get("/health_check")] +async fn health_check() -> HttpResponse { + HttpResponse::Ok().finish() +} + +#[post("/v1/embeddings")] +async fn create_embeddings(req: web::Json) -> HttpResponse { + // Validate input + if req.input.is_empty() { + return HttpResponse::BadRequest().json(ErrorResponse { + error: ErrorDetail { + message: "Input cannot be empty".to_string(), + error_type: "invalid_request_error".to_string(), + code: Some("empty_input".to_string()), + }, + }); + } + + // Create embedder + let embedder = match EmbedderBuilder::new() + .model_id(Some(req.model.as_str())) + .from_pretrained_hf() + { + Ok(embedder) => embedder, + Err(e) => { + return HttpResponse::InternalServerError().json(ErrorResponse { + error: ErrorDetail { + message: format!("Failed to initialize embedder: {}", e), + error_type: "server_error".to_string(), + code: Some("embedder_init_failed".to_string()), + }, + }); + } + }; + + // Convert input to string slices + let input_slices: Vec<&str> = req.input.iter().map(|s| s.as_str()).collect(); + + // Generate embeddings + let config = TextEmbedConfig::default(); + let embeddings = match embed_query(&input_slices, &embedder, Some(&config)).await { + Ok(embeddings) => embeddings, + Err(e) => { + return HttpResponse::InternalServerError().json(ErrorResponse { + error: ErrorDetail { + message: format!("Failed to generate embeddings: {}", e), + error_type: "server_error".to_string(), + code: Some("embedding_generation_failed".to_string()), + }, + }); + } + }; + + // Convert to OpenAI format + let embedding_data: Vec = embeddings + .into_iter() + .enumerate() + .map(|(index, embed_data)| { + let embedding_vector = match embed_data.embedding { + embed_anything::embeddings::embed::EmbeddingResult::DenseVector(vec) => vec, + embed_anything::embeddings::embed::EmbeddingResult::MultiVector(_) => { + // For multi-vector embeddings, we'll flatten them (this might need adjustment based on requirements) + vec![0.0] // Placeholder - you might want to handle this differently + } + }; + + EmbeddingData { + object: "embedding".to_string(), + index, + embedding: embedding_vector, + } + }) + .collect(); + + // Calculate usage (simplified - you might want to implement proper token counting) + let total_tokens = req.input.iter().map(|s| s.split_whitespace().count()).sum(); + + let response = OpenAIEmbedResponse { + object: "list".to_string(), + data: embedding_data, + model: req.model.clone(), + usage: Usage { + prompt_tokens: total_tokens, + total_tokens, + }, + }; + + HttpResponse::Ok().json(response) +} + +pub fn run(listener: TcpListener) -> std::io::Result { + let server = HttpServer::new(|| { + App::new() + .service(health_check) + .service(create_embeddings) + }) + .listen(listener)? + .run(); + Ok(server) +} \ No newline at end of file diff --git a/server/src/main.rs b/server/src/main.rs new file mode 100644 index 00000000..2e8db9c4 --- /dev/null +++ b/server/src/main.rs @@ -0,0 +1,9 @@ +use std::net::TcpListener; +use server::run; + +#[actix_web::main] +async fn main() -> std::io::Result<()> { + let listener = TcpListener::bind("0.0.0.0:8080")?; + println!("Listening on {}", listener.local_addr().unwrap()); + run(listener).expect("Failed to bind address").await +} \ No newline at end of file diff --git a/test_files/clip/cat3.jpg b/test_files/clip/cat3.jpg new file mode 100644 index 00000000..e131e8ec Binary files /dev/null and b/test_files/clip/cat3.jpg differ