Derived pooled fp32 and fused fp16 CUDA graphs from jinaai/jina-embeddings-v2-base-code@516f4ba
Browse files- README.md +48 -0
- onnx/model.onnx +3 -0
- onnx/model_cuda_fp16.onnx +3 -0
- 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
|
|
|