IA0x00 commited on
Commit
732cbcc
·
verified ·
1 Parent(s): c57ba24

Derived pooled fp32 and fused fp16 CUDA graphs from jinaai/jina-embeddings-v2-base-code@516f4ba

Browse files
Files changed (4) hide show
  1. README.md +48 -0
  2. onnx/model.onnx +3 -0
  3. onnx/model_cuda_fp16.onnx +3 -0
  4. tokenizer.json +0 -0
README.md ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ base_model: jinaai/jina-embeddings-v2-base-code
4
+ library_name: onnx
5
+ tags:
6
+ - feature-extraction
7
+ - sentence-similarity
8
+ - onnx
9
+ - code
10
+ ---
11
+
12
+ # jina-embeddings-v2-base-code, inference-optimized ONNX graphs for CodeSage
13
+
14
+ These are derived ONNX graphs of [`jinaai/jina-embeddings-v2-base-code`](https://huggingface.co/jinaai/jina-embeddings-v2-base-code) at revision `516f4baf13dec4ddddda8631e019b5737c8bc250`, built for the [CodeSage](https://github.com/iliaal/codesage) indexer. The weights are unchanged apart from precision. All credit for the model goes to Jina AI; the license is the upstream Apache-2.0.
15
+
16
+ ## Files
17
+
18
+ | File | Precision | Attention | Output |
19
+ |---|---|---|---|
20
+ | `onnx/model.onnx` | fp32 | upstream (decomposed) | `sentence_embedding` [batch, 768], mean-pooled and L2-normalized |
21
+ | `onnx/model_cuda_fp16.onnx` | fp16 weights, int64/fp32 I/O | `com.microsoft` `MultiHeadAttention` with ALiBi as `attention_bias` | same |
22
+ | `tokenizer.json` | | | copied unchanged from upstream |
23
+
24
+ Inputs for both graphs are `input_ids` and `attention_mask` (int64, [batch, seq]). Mean pooling over the attention mask and L2 normalization are part of the graph, so you don't pool the output yourself.
25
+
26
+ `model_cuda_fp16.onnx` targets the ONNX Runtime CUDA execution provider. Pad sequences to a multiple of 8 tokens so the memory-efficient attention kernel accepts the bias.
27
+
28
+ ## Measured effect
29
+
30
+ Measured with ONNX Runtime 1.24.4 on an RTX 4080 Laptop GPU, using 2,048 real chunks of up to 512 tokens. Cosine and neighbour overlap are against the upstream fp32 graph with host-side mean pooling.
31
+
32
+ | Corpus | Upstream fp32 | `model.onnx` | `model_cuda_fp16.onnx` | min cosine (fp16) | top-10 neighbour overlap (fp16) |
33
+ |---|---|---|---|---|---|
34
+ | php-src (C) | 59.7 chunks/s | 89 | 249–257 | 0.99999 | 0.994 |
35
+ | Laravel app (PHP) | 67.2 | 113 | 301–317 | 0.999996 | 0.996 |
36
+ | React app (TS/JS) | 69.8 | 109 | 296–317 | 0.999995 | 0.997 |
37
+
38
+ `model.onnx` matches the upstream vectors to within float rounding (min cosine 0.999999).
39
+
40
+ ## Reproduce
41
+
42
+ The script `scripts/derive-jina-onnx.py` in the CodeSage repository downloads the pinned upstream files, verifies their sha256, and writes these files byte-for-byte. It was run with onnx 1.21.0 and onnxruntime 1.24.4.
43
+
44
+ ```
45
+ tokenizer.json b01c78a902aa4facb2f47f95449f48e2f7bbfea5d2472ee2f6ce92323c6f86e5
46
+ onnx/model.onnx 98ca3fc0def59e9fea861ca8888f805fd3cd10d429695ff53bd03ab053e8aee4
47
+ onnx/model_cuda_fp16.onnx 5d34291e0d44a01924d6709c631b141a875a7e9add7f118dde87050644cfc2ab
48
+ ```
onnx/model.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:98ca3fc0def59e9fea861ca8888f805fd3cd10d429695ff53bd03ab053e8aee4
3
+ size 641517835
onnx/model_cuda_fp16.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5d34291e0d44a01924d6709c631b141a875a7e9add7f118dde87050644cfc2ab
3
+ size 320904215
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff