jina-v2-code-query-4l

A 4-layer distillation of jina-embeddings-v2-base-code for embedding search queries, as int8 ONNX (85 MB). It keeps the base model's token embeddings and 4 of its 12 layers (0, 4, 8 and 11), trained to point where the base model points. Made for provgrep, which orders code embedded by the base model by how close it is to a query embedded by this one: a query costs a third of the CPU.

It is for queries, not for code. Embed documents with the base model.

Use

• Tokenizer: the base model's (tokenizer.json, the same file), truncated at 512 tokens.
• Inputs input_ids and attention_mask; output last_hidden_state.
• Pool by the mean over the attention mask, then normalize, as for the base model.

How it was made

• Texts: 5,111,554 query-like texts from the 10,915 GitHub repositories of provgrep's corpus (corpus/github.json, sha256 4da3293d763e): first sentences of functions' comments and docstrings, function names split into words, a sentence with an identifier its function calls, and the first words of a sentence, as typed. provgrep's evaluation questions are not among them.
• Targets: the base model's embedding of each text (fp16, revision 516f4baf13de).
• Loss: one minus the cosine of the two, after projecting both on the 256 directions provgrep's code embeddings vary most along (a PCA of 32.5 million of them). Closeness elsewhere was not trained for.
• One epoch over the 5,082,756 texts of at most 64 tokens: batches of 256, AdamW at 1e-4 with a one-cycle schedule, bf16, about 24 minutes on an RTX A4000. Exported at opset 17 and quantized with onnxruntime's dynamic int8 quantization.

How it does

• Held out (20,000 texts): cosine 0.970 with the base model in that projected space.
• provgrep's four suites of search intents (provgrep evaluate search with the code of provgrep 5904e2f, 163 intents asked with any license and with permissive ones; on its index of corpus sha256 4da3293d763e, results reordered by cosine with each): a hit first for 139 searches and in the top 10 for 263, against 134 and 259 with the base model in int8.
• On an 8-core server, one query of those suites: 18.6 CPU-ms on one thread, against 59 for the base model in int8; 6.6 ms on four threads.

License

Apache-2.0, as the base model, whose weights are Jina AI's.