wip

Co-Authored-By: Antonio <antonio@zed.dev>

Cargo.lock

@@ -550,7 +550,7 @@ dependencies = [
  "libc",
  "pin-project",
  "redox_syscall 0.2.16",
- "xattr",
+ "xattr 0.2.3",
 ]
 
 [[package]]
@@ -974,7 +974,7 @@ dependencies = [
  "collections",
  "editor",
  "gpui",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "outline",
  "project",
@@ -1924,7 +1924,7 @@ dependencies = [
  "cranelift-codegen",
  "cranelift-entity",
  "cranelift-frontend",
- "itertools",
+ "itertools 0.10.5",
  "log",
  "smallvec",
  "wasmparser",
@@ -2008,6 +2008,12 @@ dependencies = [
  "cfg-if 1.0.0",
 ]
 
+[[package]]
+name = "crunchy"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
+
 [[package]]
 name = "crypto-common"
 version = "0.1.6"
@@ -2068,6 +2074,41 @@ dependencies = [
  "winapi 0.3.9",
 ]
 
+[[package]]
+name = "darling"
+version = "0.14.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b750cb3417fd1b327431a470f388520309479ab0bf5e323505daf0290cd3850"
+dependencies = [
+ "darling_core",
+ "darling_macro",
+]
+
+[[package]]
+name = "darling_core"
+version = "0.14.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "109c1ca6e6b7f82cc233a97004ea8ed7ca123a9af07a8230878fcfda9b158bf0"
+dependencies = [
+ "fnv",
+ "ident_case",
+ "proc-macro2",
+ "quote",
+ "strsim",
+ "syn 1.0.109",
+]
+
+[[package]]
+name = "darling_macro"
+version = "0.14.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4aab4dbc9f7611d8b55048a3a16d2d010c2c8334e46304b40ac1cc14bf3b48e"
+dependencies = [
+ "darling_core",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "dashmap"
 version = "5.5.1"
@@ -2137,6 +2178,37 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "derive_builder"
+version = "0.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d67778784b508018359cbc8696edb3db78160bab2c2a28ba7f56ef6932997f8"
+dependencies = [
+ "derive_builder_macro",
+]
+
+[[package]]
+name = "derive_builder_core"
+version = "0.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c11bdc11a0c47bc7d37d582b5285da6849c96681023680b906673c5707af7b0f"
+dependencies = [
+ "darling",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
+[[package]]
+name = "derive_builder_macro"
+version = "0.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ebcda35c7a396850a55ffeac740804b40ffec779b98fffbb1738f4033f0ee79e"
+dependencies = [
+ "derive_builder_core",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "derive_more"
 version = "0.99.17"
@@ -2345,7 +2417,7 @@ dependencies = [
  "git",
  "gpui",
  "indoc",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "lazy_static",
  "log",
@@ -2474,6 +2546,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "esaxx-rs"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f748b253ceca9fed5f42f8b5ceb3851e93102199bc25b64b65369f76e5c0a35"
+
 [[package]]
 name = "etagere"
 version = "0.2.8"
@@ -3135,7 +3213,7 @@ dependencies = [
  "futures 0.3.28",
  "gpui_macros",
  "image",
- "itertools",
+ "itertools 0.10.5",
  "lazy_static",
  "log",
  "media",
@@ -3206,6 +3284,16 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "half"
+version = "2.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc52e53916c08643f1b56ec082790d1e86a32e58dc5268f897f313fbae7b4872"
+dependencies = [
+ "cfg-if 1.0.0",
+ "crunchy",
+]
+
 [[package]]
 name = "hashbrown"
 version = "0.11.2"
@@ -3501,6 +3589,12 @@ dependencies = [
  "cc",
 ]
 
+[[package]]
+name = "ident_case"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
+
 [[package]]
 name = "idna"
 version = "0.4.0"
@@ -3709,6 +3803,24 @@ dependencies = [
  "waker-fn",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f56a2d0bc861f9165be4eb3442afd3c236d8a98afd426f65d92324ae1091a484"
+dependencies = [
+ "either",
+]
+
+[[package]]
+name = "itertools"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "284f18f85651fe11e8a991b2adb42cb078325c996ed026d994719efcfca1d54b"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itertools"
 version = "0.10.5"
@@ -4205,6 +4317,22 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "macro_rules_attribute"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf0c9b980bf4f3a37fd7b1c066941dd1b1d0152ce6ee6e8fe8c49b9f6810d862"
+dependencies = [
+ "macro_rules_attribute-proc_macro",
+ "paste",
+]
+
+[[package]]
+name = "macro_rules_attribute-proc_macro"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "58093314a45e00c77d5c508f76e77c3396afbbc0d01506e7fae47b018bac2b1d"
+
 [[package]]
 name = "malloc_buf"
 version = "0.0.6"
@@ -4474,6 +4602,27 @@ dependencies = [
  "winapi 0.3.9",
 ]
 
+[[package]]
+name = "monostate"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "15f370ae88093ec6b11a710dec51321a61d420fafd1bad6e30d01bd9c920e8ee"
+dependencies = [
+ "monostate-impl",
+ "serde",
+]
+
+[[package]]
+name = "monostate-impl"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "371717c0a5543d6a800cac822eac735aa7d2d2fbb41002e9856a4089532dbdce"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.29",
+]
+
 [[package]]
 name = "more-asserts"
 version = "0.2.2"
@@ -4513,6 +4662,19 @@ dependencies = [
  "tempfile",
 ]
 
+[[package]]
+name = "ndarray"
+version = "0.15.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32"
+dependencies = [
+ "matrixmultiply",
+ "num-complex 0.4.4",
+ "num-integer",
+ "num-traits",
+ "rawpointer",
+]
+
 [[package]]
 name = "ndk"
 version = "0.7.0"
@@ -4640,7 +4802,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b8536030f9fea7127f841b45bb6243b27255787fb4eb83958aa1ef9d2fdc0c36"
 dependencies = [
  "num-bigint 0.2.6",
- "num-complex",
+ "num-complex 0.2.4",
  "num-integer",
  "num-iter",
  "num-rational 0.2.4",
@@ -4696,6 +4858,15 @@ dependencies = [
  "num-traits",
 ]
 
+[[package]]
+name = "num-complex"
+version = "0.4.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "num-derive"
 version = "0.3.3"
@@ -4876,6 +5047,28 @@ version = "1.18.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d"
 
+[[package]]
+name = "onig"
+version = "6.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c4b31c8722ad9171c6d77d3557db078cab2bd50afcc9d09c8b315c59df8ca4f"
+dependencies = [
+ "bitflags 1.3.2",
+ "libc",
+ "once_cell",
+ "onig_sys",
+]
+
+[[package]]
+name = "onig_sys"
+version = "69.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b829e3d7e9cc74c7e315ee8edb185bf4190da5acde74afd7fc59c35b1f086e7"
+dependencies = [
+ "cc",
+ "pkg-config",
+]
+
 [[package]]
 name = "opaque-debug"
 version = "0.3.0"
@@ -4935,6 +5128,26 @@ dependencies = [
  "num-traits",
 ]
 
+[[package]]
+name = "ort"
+version = "1.15.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5e56c9c4185ee949ef961aca8777d1dbd52cb104b444669adad63e8181820a7"
+dependencies = [
+ "flate2",
+ "half",
+ "lazy_static",
+ "libc",
+ "ndarray",
+ "tar",
+ "thiserror",
+ "tracing",
+ "ureq",
+ "vswhom",
+ "winapi 0.3.9",
+ "zip",
+]
+
 [[package]]
 name = "os_str_bytes"
 version = "6.5.1"
@@ -5474,7 +5687,7 @@ dependencies = [
  "globset",
  "gpui",
  "ignore",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "lazy_static",
  "log",
@@ -5598,7 +5811,7 @@ checksum = "62941722fb675d463659e49c4f3fe1fe792ff24fe5bbaa9c08cd3b98a1c354f5"
 dependencies = [
  "bytes 1.4.0",
  "heck 0.3.3",
- "itertools",
+ "itertools 0.10.5",
  "lazy_static",
  "log",
  "multimap",
@@ -5617,7 +5830,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "600d2f334aa05acb02a755e217ef1ab6dea4d51b58b7846588b747edec04efba"
 dependencies = [
  "anyhow",
- "itertools",
+ "itertools 0.10.5",
  "proc-macro2",
  "quote",
  "syn 1.0.109",
@@ -5630,7 +5843,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9cc1a3263e07e0bf68e96268f37665207b49560d98739662cdfaae215c720fe"
 dependencies = [
  "anyhow",
- "itertools",
+ "itertools 0.10.5",
  "proc-macro2",
  "quote",
  "syn 1.0.109",
@@ -5859,6 +6072,17 @@ dependencies = [
  "rayon-core",
 ]
 
+[[package]]
+name = "rayon-cond"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd1259362c9065e5ea39a789ef40b1e3fd934c94beb7b5ab3ac6629d3b5e7cb7"
+dependencies = [
+ "either",
+ "itertools 0.8.2",
+ "rayon",
+]
+
 [[package]]
 name = "rayon-core"
 version = "1.11.0"
@@ -6404,6 +6628,18 @@ dependencies = [
  "webpki 0.22.0",
 ]
 
+[[package]]
+name = "rustls"
+version = "0.21.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd8d6c9f025a446bc4d18ad9632e69aec8f287aa84499ee335599fabd20c3fd8"
+dependencies = [
+ "log",
+ "ring",
+ "rustls-webpki 0.101.4",
+ "sct 0.7.0",
+]
+
 [[package]]
 name = "rustls-pemfile"
 version = "1.0.3"
@@ -6413,6 +6649,26 @@ dependencies = [
  "base64 0.21.2",
 ]
 
+[[package]]
+name = "rustls-webpki"
+version = "0.100.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e98ff011474fa39949b7e5c0428f9b4937eda7da7848bbb947786b7be0b27dab"
+dependencies = [
+ "ring",
+ "untrusted",
+]
+
+[[package]]
+name = "rustls-webpki"
+version = "0.101.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7d93931baf2d282fff8d3a532bbfd7653f734643161b87e3e01e59a04439bf0d"
+dependencies = [
+ "ring",
+ "untrusted",
+]
+
 [[package]]
 name = "rustversion"
 version = "1.0.14"
@@ -6734,6 +6990,8 @@ dependencies = [
  "lazy_static",
  "log",
  "matrixmultiply",
+ "ndarray",
+ "ort",
  "parking_lot 0.11.2",
  "parse_duration",
  "picker",
@@ -6752,6 +7010,7 @@ dependencies = [
  "tempdir",
  "theme",
  "tiktoken-rs 0.5.1",
+ "tokenizers",
  "tree-sitter",
  "tree-sitter-cpp",
  "tree-sitter-elixir",
@@ -7209,6 +7468,18 @@ dependencies = [
  "lock_api",
 ]
 
+[[package]]
+name = "spm_precompiled"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
+dependencies = [
+ "base64 0.13.1",
+ "nom",
+ "serde",
+ "unicode-segmentation",
+]
+
 [[package]]
 name = "spsc-buffer"
 version = "0.1.1"
@@ -7248,7 +7519,7 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0c12bc9199d1db8234678b7051747c07f517cdcf019262d1847b94ec8b1aee3e"
 dependencies = [
- "itertools",
+ "itertools 0.10.5",
  "nom",
  "unicode_categories",
 ]
@@ -7586,6 +7857,17 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
 
+[[package]]
+name = "tar"
+version = "0.4.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b16afcea1f22891c49a00c751c7b63b2233284064f11a200fc624137c51e2ddb"
+dependencies = [
+ "filetime",
+ "libc",
+ "xattr 1.0.1",
+]
+
 [[package]]
 name = "target-lexicon"
 version = "0.12.11"
@@ -7634,7 +7916,7 @@ dependencies = [
  "dirs 4.0.0",
  "futures 0.3.28",
  "gpui",
- "itertools",
+ "itertools 0.10.5",
  "lazy_static",
  "libc",
  "mio-extras",
@@ -7665,7 +7947,7 @@ dependencies = [
  "editor",
  "futures 0.3.28",
  "gpui",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "lazy_static",
  "libc",
@@ -7913,6 +8195,37 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
+[[package]]
+name = "tokenizers"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "12b515a66453a4d68f03398054f7204fd0dde6b93d3f20ea90b08025ab49b499"
+dependencies = [
+ "aho-corasick 0.7.20",
+ "derive_builder",
+ "esaxx-rs",
+ "getrandom 0.2.10",
+ "itertools 0.9.0",
+ "lazy_static",
+ "log",
+ "macro_rules_attribute",
+ "monostate",
+ "onig",
+ "paste",
+ "rand 0.8.5",
+ "rayon",
+ "rayon-cond",
+ "regex",
+ "regex-syntax 0.7.4",
+ "serde",
+ "serde_json",
+ "spm_precompiled",
+ "thiserror",
+ "unicode-normalization-alignments",
+ "unicode-segmentation",
+ "unicode_categories",
+]
+
 [[package]]
 name = "tokio"
 version = "1.32.0"
@@ -8615,6 +8928,15 @@ dependencies = [
  "tinyvec",
 ]
 
+[[package]]
+name = "unicode-normalization-alignments"
+version = "0.1.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
+dependencies = [
+ "smallvec",
+]
+
 [[package]]
 name = "unicode-script"
 version = "0.5.5"
@@ -8657,6 +8979,21 @@ version = "0.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a"
 
+[[package]]
+name = "ureq"
+version = "2.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b11c96ac7ee530603dcdf68ed1557050f374ce55a5a07193ebf8cbc9f8927e9"
+dependencies = [
+ "base64 0.21.2",
+ "log",
+ "once_cell",
+ "rustls 0.21.7",
+ "rustls-webpki 0.100.2",
+ "url",
+ "webpki-roots 0.23.1",
+]
+
 [[package]]
 name = "url"
 version = "2.4.0"
@@ -8831,7 +9168,7 @@ dependencies = [
  "editor",
  "gpui",
  "indoc",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "language_selector",
  "log",
@@ -8849,6 +9186,26 @@ dependencies = [
  "workspace",
 ]
 
+[[package]]
+name = "vswhom"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be979b7f07507105799e854203b470ff7c78a1639e330a58f183b5fea574608b"
+dependencies = [
+ "libc",
+ "vswhom-sys",
+]
+
+[[package]]
+name = "vswhom-sys"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d3b17ae1f6c8a2b28506cd96d412eebf83b4a0ff2cbefeeb952f2f9dfa44ba18"
+dependencies = [
+ "cc",
+ "libc",
+]
+
 [[package]]
 name = "vte"
 version = "0.11.1"
@@ -9315,6 +9672,15 @@ dependencies = [
  "webpki 0.22.0",
 ]
 
+[[package]]
+name = "webpki-roots"
+version = "0.23.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b03058f88386e5ff5310d9111d53f48b17d732b401aeb83a8d5190f2ac459338"
+dependencies = [
+ "rustls-webpki 0.100.2",
+]
+
 [[package]]
 name = "weezl"
 version = "0.1.7"
@@ -9663,7 +10029,7 @@ dependencies = [
  "gpui",
  "indoc",
  "install_cli",
- "itertools",
+ "itertools 0.10.5",
  "language",
  "lazy_static",
  "log",
@@ -9711,6 +10077,15 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "xattr"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4686009f71ff3e5c4dbcf1a282d0a44db3f021ba69350cd42086b3e5f1c6985"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "xmlparser"
 version = "0.13.5"
@@ -9914,6 +10289,18 @@ dependencies = [
  "syn 2.0.29",
 ]
 
+[[package]]
+name = "zip"
+version = "0.6.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261"
+dependencies = [
+ "byteorder",
+ "crc32fast",
+ "crossbeam-utils",
+ "flate2",
+]
+
 [[package]]
 name = "zstd"
 version = "0.11.2+zstd.1.5.2"

crates/semantic_index/Cargo.toml

@@ -41,6 +41,9 @@ schemars.workspace = true
 globset.workspace = true
 sha1 = "0.10.5"
 parse_duration = "2.1.1"
+ort = { version = "1.15.2", features = ["coreml"]}
+tokenizers = { version = ">=0.13.4", default-features = false, features = [ "onig" ] }
+ndarray = "0.15"
 
 [dev-dependencies]
 collections = { path = "../collections", features = ["test-support"] }

crates/semantic_index/src/cross_encoder.rs

@@ -0,0 +1,115 @@
+use ndarray::CowArray;
+use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, Session, SessionBuilder, Value};
+use tokenizers::Tokenizer;
+use util::paths::MODELS_DIR;
+
+pub struct CrossEncoder {
+    session: Session,
+    tokenizer: Tokenizer,
+}
+
+fn sigmoid(val: f32) -> f32 {
+    1.0 / (1.0 + (-val).exp())
+}
+
+impl CrossEncoder {
+    pub fn load() -> anyhow::Result<Self> {
+        let model_path = MODELS_DIR.join("cross-encoder").join("model.onnx");
+        let tokenizer_path = MODELS_DIR.join("cross-encoder").join("tokenizer.json");
+
+        let environment = Environment::builder()
+            .with_name("cross-encoder")
+            .with_execution_providers([ExecutionProvider::CoreML(Default::default())])
+            .build()?
+            .into_arc();
+
+        let session = SessionBuilder::new(&environment)?
+            .with_optimization_level(GraphOptimizationLevel::Level1)?
+            .with_model_from_file(model_path)?;
+
+        let mut tokenizer = Tokenizer::from_file(tokenizer_path).unwrap();
+        tokenizer
+            .with_truncation(Some(tokenizers::TruncationParams {
+                direction: Default::default(),
+                max_length: 512,
+                strategy: Default::default(),
+                stride: 0,
+            }))
+            .unwrap();
+
+        Ok(Self { session, tokenizer })
+    }
+
+    pub fn score(&self, query: &str, candidates: &[String]) -> anyhow::Result<Vec<f32>> {
+        let spans = candidates
+            .into_iter()
+            .map(|candidate| format!("{}. {}", query, candidate))
+            .collect::<Vec<_>>();
+
+        let encodings = self.tokenizer.encode_batch(spans, true).unwrap();
+
+        let mut results = Vec::new();
+        for encoding in encodings {
+            // Get Input Variables Individually
+            let input_ids = encoding.get_ids();
+            let attention_mask = encoding.get_attention_mask();
+            let token_type_ids = encoding.get_type_ids();
+            let length = input_ids.len();
+
+            // Convert to Arrays
+            let inputs_ids_array = CowArray::from(ndarray::Array::from_shape_vec(
+                (1, length),
+                input_ids.iter().map(|&x| x as i64).collect(),
+            )?);
+
+            let attention_mask_array = CowArray::from(ndarray::Array::from_shape_vec(
+                (1, length),
+                attention_mask.iter().map(|&x| x as i64).collect(),
+            )?)
+            .into_dyn();
+
+            let token_type_ids_array = CowArray::from(ndarray::Array::from_shape_vec(
+                (1, length),
+                token_type_ids.iter().map(|&x| x as i64).collect(),
+            )?)
+            .into_dyn();
+
+            let outputs = self.session.run(vec![
+                Value::from_array(self.session.allocator(), &inputs_ids_array.into_dyn())?,
+                Value::from_array(self.session.allocator(), &attention_mask_array)?,
+                Value::from_array(self.session.allocator(), &token_type_ids_array)?,
+            ]);
+
+            let output = outputs.unwrap()[0].try_extract::<f32>().unwrap();
+            let value = output.view().to_owned();
+
+            let val = value.as_slice().unwrap()[0];
+            results.push(sigmoid(val))
+        }
+
+        Ok(results)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_cross_encoder() {
+        let cross_encoder = CrossEncoder::load().unwrap();
+        let results = cross_encoder
+            .score(
+                "I like you",
+                &[
+                    "I hate you.".into(),
+                    "I love you.".into(),
+                    "my name is kyle".into(),
+                ],
+            )
+            .unwrap();
+        assert_eq!(results.len(), 3);
+        assert!(results[1] > results[0]);
+        assert!(results[0] > results[2]);
+    }
+}

crates/semantic_index/src/semantic_index.rs

@@ -1,3 +1,4 @@
+mod cross_encoder;
 mod db;
 mod embedding;
 mod embedding_queue;
@@ -7,7 +8,7 @@ pub mod semantic_index_settings;
 #[cfg(test)]
 mod semantic_index_tests;
 
-use crate::semantic_index_settings::SemanticIndexSettings;
+use crate::{cross_encoder::CrossEncoder, semantic_index_settings::SemanticIndexSettings};
 use anyhow::{anyhow, Result};
 use collections::{BTreeMap, HashMap, HashSet};
 use db::VectorDatabase;
@@ -265,6 +266,7 @@ pub struct PendingFile {
 pub struct SearchResult {
     pub buffer: ModelHandle<Buffer>,
     pub range: Range<Anchor>,
+    pub similarity: f32,
 }
 
 impl SemanticIndex {
@@ -696,7 +698,7 @@ impl SemanticIndex {
         let embedding_provider = self.embedding_provider.clone();
         let db_path = self.db.path().clone();
         let fs = self.fs.clone();
-        cx.spawn(|this, mut cx| async move {
+        cx.spawn(|this, cx| async move {
             index.await?;
 
             let t0 = Instant::now();
@@ -708,7 +710,7 @@ impl SemanticIndex {
             }
 
             let phrase_embedding = embedding_provider
-                .embed_batch(vec![phrase])
+                .embed_batch(vec![phrase.clone()])
                 .await?
                 .into_iter()
                 .next()
@@ -749,6 +751,11 @@ impl SemanticIndex {
                 ids_len / batch_n
             };
 
+            let cross_encoder = Arc::new(
+                cx.background()
+                    .spawn(async move { CrossEncoder::load() })
+                    .await?,
+            );
             let mut batch_results = Vec::new();
             for batch in file_ids.chunks(batch_size) {
                 let batch = batch.into_iter().map(|v| *v).collect::<Vec<i64>>();
@@ -756,77 +763,126 @@ impl SemanticIndex {
                 let fs = fs.clone();
                 let db_path = db_path.clone();
                 let phrase_embedding = phrase_embedding.clone();
+                let phrase = phrase.clone();
+                let cross_encoder = cross_encoder.clone();
+                let project = project.clone();
                 if let Some(db) = VectorDatabase::new(fs, db_path.clone(), cx.background())
                     .await
                     .log_err()
                 {
-                    batch_results.push(async move {
-                        db.top_k_search(&phrase_embedding, limit, batch.as_slice())
-                            .await
-                    });
+                    let this = this.clone();
+                    batch_results.push(cx.spawn(|mut cx| async move {
+                        let span_ids = db
+                            .top_k_search(&phrase_embedding, limit, batch.as_slice())
+                            .await?
+                            .into_iter()
+                            .map(|(span_id, _)| span_id)
+                            .collect::<Vec<_>>();
+                        let mut spans_by_buffer = HashMap::default();
+                        for (worktree_db_id, path, range) in db.spans_for_ids(&span_ids).await? {
+                            let worktree_id = this.read_with(&cx, |this, _| {
+                                let project_state = this
+                                    .projects
+                                    .get(&project.downgrade())
+                                    .ok_or_else(|| anyhow!("project not added"))?;
+
+                                anyhow::Ok(project_state.worktree_id_for_db_id(worktree_db_id))
+                            })?;
+
+                            if let Some(worktree_id) = worktree_id {
+                                let buffer = project
+                                    .update(&mut cx, |project, cx| {
+                                        project.open_buffer((worktree_id, path), cx)
+                                    })
+                                    .await
+                                    .log_err();
+                                if let Some(buffer) = buffer {
+                                    let range = buffer.read_with(&cx, |buffer, _| {
+                                        let range = buffer.clip_offset(range.start, Bias::Left)
+                                            ..buffer.clip_offset(range.end, Bias::Right);
+                                        buffer.anchor_before(range.start)
+                                            ..buffer.anchor_after(range.end)
+                                    });
+
+                                    spans_by_buffer
+                                        .entry(buffer)
+                                        .or_insert(Vec::new())
+                                        .push(range);
+                                }
+                            }
+                        }
+
+                        let mut spans = Vec::new();
+                        for (buffer, ranges) in &spans_by_buffer {
+                            buffer.read_with(&cx, |buffer, _| {
+                                for range in ranges {
+                                    let span =
+                                        buffer.text_for_range(range.clone()).collect::<String>();
+                                    spans.push(span);
+                                }
+                            });
+                        }
+
+                        // Cross Encoder
+                        // TODO: move background.spawn into cross_encoder.
+                        let results = cx
+                            .background()
+                            .spawn(async move {
+                                let mut results = Vec::new();
+                                let mut scores = cross_encoder.score(&phrase, &spans)?.into_iter();
+                                for (buffer, ranges) in spans_by_buffer {
+                                    for range in ranges {
+                                        let similarity = if let Some(similarity) = scores.next() {
+                                            similarity
+                                        } else {
+                                            log::error!("cross encoder returned too few scores");
+                                            f32::NEG_INFINITY
+                                        };
+
+                                        results.push(SearchResult {
+                                            buffer: buffer.clone(),
+                                            range,
+                                            similarity,
+                                        });
+                                    }
+                                }
+                                anyhow::Ok(results)
+                            })
+                            .await?;
+
+                        anyhow::Ok(results)
+                    }));
                 }
             }
             let batch_results = futures::future::join_all(batch_results).await;
 
-            let mut results = Vec::new();
+            let mut results = Vec::<SearchResult>::new();
             for batch_result in batch_results {
-                if batch_result.is_ok() {
-                    for (id, similarity) in batch_result.unwrap() {
-                        let ix = match results.binary_search_by(|(_, s)| {
-                            similarity.partial_cmp(&s).unwrap_or(Ordering::Equal)
+                if let Some(batch_result) = batch_result.log_err() {
+                    for new_result in batch_result {
+                        let ix = match results.binary_search_by(|old_result| {
+                            new_result
+                                .similarity
+                                .partial_cmp(&old_result.similarity)
+                                .unwrap_or(Ordering::Equal)
                         }) {
                             Ok(ix) => ix,
                             Err(ix) => ix,
                         };
-                        results.insert(ix, (id, similarity));
+                        dbg!(ix);
+                        dbg!(new_result.similarity);
+                        results.insert(ix, new_result);
                         results.truncate(limit);
                     }
                 }
             }
 
-            let ids = results.into_iter().map(|(id, _)| id).collect::<Vec<i64>>();
-            let spans = database.spans_for_ids(ids.as_slice()).await?;
-
-            let mut tasks = Vec::new();
-            let mut ranges = Vec::new();
-            let weak_project = project.downgrade();
-            project.update(&mut cx, |project, cx| {
-                for (worktree_db_id, file_path, byte_range) in spans {
-                    let project_state =
-                        if let Some(state) = this.read(cx).projects.get(&weak_project) {
-                            state
-                        } else {
-                            return Err(anyhow!("project not added"));
-                        };
-                    if let Some(worktree_id) = project_state.worktree_id_for_db_id(worktree_db_id) {
-                        tasks.push(project.open_buffer((worktree_id, file_path), cx));
-                        ranges.push(byte_range);
-                    }
-                }
-
-                Ok(())
-            })?;
-
-            let buffers = futures::future::join_all(tasks).await;
-
             log::trace!(
                 "Semantic Searching took: {:?} milliseconds in total",
                 t0.elapsed().as_millis()
             );
 
-            Ok(buffers
-                .into_iter()
-                .zip(ranges)
-                .filter_map(|(buffer, range)| {
-                    let buffer = buffer.log_err()?;
-                    let range = buffer.read_with(&cx, |buffer, _| {
-                        let start = buffer.clip_offset(range.start, Bias::Left);
-                        let end = buffer.clip_offset(range.end, Bias::Right);
-                        buffer.anchor_before(start)..buffer.anchor_after(end)
-                    });
-                    Some(SearchResult { buffer, range })
-                })
-                .collect::<Vec<_>>())
+            Ok(results)
         })
     }
 

crates/util/src/paths.rs

@@ -18,6 +18,7 @@ lazy_static::lazy_static! {
     pub static ref LOG: PathBuf = LOGS_DIR.join("Zed.log");
     pub static ref OLD_LOG: PathBuf = LOGS_DIR.join("Zed.log.old");
     pub static ref LOCAL_SETTINGS_RELATIVE_PATH: &'static Path = Path::new(".zed/settings.json");
+    pub static ref MODELS_DIR: PathBuf = HOME.join("Library/Application Support/Zed/models");
 }
 
 pub mod legacy {