onnx-checkpoint

.gitattributes

@@ -34,3 +34,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+*.original filter=lfs diff=lfs merge=lfs -text
+onnx/roberta.embeddings.word_embeddings.parametrizations.weight.0.lora_A filter=lfs diff=lfs merge=lfs -text
+onnx/_roberta_encoder_layers.0_mixer_rotary_emb_Constant_6_attr__value filter=lfs diff=lfs merge=lfs -text
+onnx/_roberta_encoder_layers.0_mixer_rotary_emb_Constant_attr__value filter=lfs diff=lfs merge=lfs -text

README.md

@@ -21528,7 +21528,7 @@ model-index:
 </p>
 
 <p align="center">
-<b>Jina Embedding V3: A Multilingual Multi-Task Embedding Model</b>
+<b>jina-embeddings-v3: Multilingual Embeddings With Task LoRA</b>
 </p>
 
 ## Quick Start
@@ -21541,12 +21541,12 @@ The easiest way to start using `jina-embeddings-v3` is with the [Jina Embedding
 
 `jina-embeddings-v3` is a **multilingual multi-task text embedding model** designed for a variety of NLP applications.
 Based on the [Jina-XLM-RoBERTa architecture](https://huggingface.co/jinaai/xlm-roberta-flash-implementation), 
-this model supports [Rotary Position Embeddings (RoPE)](https://arxiv.org/abs/2104.09864) to handle long input sequences up to **8192 tokens**.
-Additionally, it features 5 [LoRA](https://arxiv.org/abs/2106.09685) adapters to generate task-specific embeddings efficiently.
+this model supports Rotary Position Embeddings to handle long input sequences up to **8192 tokens**.
+Additionally, it features 5 LoRA adapters to generate task-specific embeddings efficiently.
 
 ### Key Features:
 - **Extended Sequence Length:** Supports up to 8192 tokens with RoPE.
-- **Task-Specific Embedding:** Customize embeddings through the `task_type` argument with the following options:
+- **Task-Specific Embedding:** Customize embeddings through the `task` argument with the following options:
     - `retrieval.query`: Used for query embeddings in asymmetric retrieval tasks
     - `retrieval.passage`: Used for passage embeddings in asymmetric retrieval tasks
     - `separation`: Used for embeddings in clustering and re-ranking applications
@@ -21560,11 +21560,6 @@ While the foundation model supports 89 languages, we've focused our tuning effor
 Hindi, Indonesian, Italian, Japanese, Korean, Latvian, Norwegian, Polish, Portuguese, Romanian, 
 Russian, Slovak, Spanish, Swedish, Thai, Turkish, Ukrainian, Urdu,** and **Vietnamese.**
 
-
-## Data & Parameters
-
-The data and training details are described in the technical report (coming soon).
-
 ## Usage
 
 **<details><summary>Apply mean pooling when integrating the model.</summary>**
@@ -21605,7 +21600,7 @@ model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code
 encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
 
 with torch.no_grad():
-    model_output = model(**encoded_input, task_type='retrieval.query')
+    model_output = model(**encoded_input, task='retrieval.query')
 
 embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
 embeddings = F.normalize(embeddings, p=2, dim=1)
@@ -21643,10 +21638,10 @@ texts = [
     "Folge dem weißen Kaninchen.",  # German
 ]
 
-# When calling the `encode` function, you can choose a `task_type` based on the use case:
+# When calling the `encode` function, you can choose a `task` based on the use case:
 # 'retrieval.query', 'retrieval.passage', 'separation', 'classification', 'text-matching'
-# Alternatively, you can choose not to pass a `task_type`, and no specific LoRA adapter will be used.
-embeddings = model.encode(texts, task_type="text-matching")
+# Alternatively, you can choose not to pass a `task`, and no specific LoRA adapter will be used.
+embeddings = model.encode(texts, task="text-matching")
 
 # Compute similarities
 print(embeddings[0] @ embeddings[1].T)
@@ -21680,61 +21675,50 @@ from sentence_transformers import SentenceTransformer
 
 model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)
 
-task_type = "retrieval.query"
+task = "retrieval.query"
 embeddings = model.encode(
     ["What is the weather like in Berlin today?"],
-    task_type=task_type,
-    prompt_name=task_type,
+    task=task,
+    prompt_name=task,
 )
 ```
 
+**<details><summary>ONNX Inference.</summary>**
+<p>
 
+You can use ONNX for efficient inference with `jina-embeddings-v3`:
+```python
+import onnxruntime
+import numpy as np
+from transformers import AutoTokenizer, PretrainedConfig
 
-## Performance
-
-### English MTEB
-| Model                          |  Dimension  |  Average  | Classification   | Clustering   | Pair Classification   | Reranking   | Retrieval   | STS  | Summarization   |
-|:------------------------------:|:-----------:|:---------:|:----:|:----:|:----:|:----:|:----:|:----:|:----:|
-| jina-embeddings-v3             |    1024     | **65.60** | **82.58**| 45.27| 84.01| 58.13| 53.87| **85.8** | 30.98|
-| jina-embeddings-v2-en          |     768     |   58.12   |   68.82    | 40.08| 84.44| 55.09| 45.64| 80.00| 30.56|
-| text-embedding-3-large         |    3072     |   62.03   |   75.45    | 49.01| 84.22| 59.16| 55.44| 81.04| 29.92|
-| multilingual-e5-large-instruct |    1024     |   64.41   |   77.56   | 47.1 | 86.19| 58.58| 52.47| 84.78| 30.39|
-| Cohere-embed-multilingual-v3.0 | 1024 |   60.08   |   64.01   | 46.6 | 86.15| 57.86| 53.84| 83.15| 30.99|
-
-### Multilingual MTEB
-
-|             Model              | Dimension |  Average  | Classification | Clustering | Pair Classification | Reranking | Retrieval |    STS    | Summarization |
-|:------------------------------:|:---------:|:---------:|:--------------:|:----------:|:-------------------:|:---------:|:---------:|:---------:|:-------------:|
-|       jina-embeddings-v3       |   1024    | **64.44** |   **71.46**    |   46.71    |        76.91        |   63.98   |   57.98   | **69.83** |       -       |
-|     multilingual-e5-large      |   1024    |   59.58   |     65.22      |   42.12    |        76.95        |   63.4    |   52.37   |   64.65   |       -       |
-| multilingual-e5-large-instruct |   1024    |   64.25   |     67.45      | **52.12**  |        77.79        | **69.02** | **58.38** |   68.77   |       -       |
-
-
-### Long Context Tasks (LongEmbed)
+# Load tokenizer and model config
+tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v3')
+config = PretrainedConfig.from_pretrained('jinaai/jina-embeddings-v3')
 
-|         Model          | Dimension |  Average  | NarrativeQA |  Needle   |  Passkey   |   QMSum   | SummScreen |  WikiQA   |
-|:----------------------:|:---------:|:---------:|:-----------:|:---------:|:----------:|:---------:|:----------:|:---------:|
-|  jina-embeddings-v3*   |   1024    | **70.39** |    33.32    | **84.00** | **100.00** | **39.75** |   92.78    |   72.46   |
-|   jina-embeddings-v2   |    768    |   58.12   |    37.89    |   54.25   |   50.25    |   38.87   |   93.48    |   73.99   |
-| text-embedding-3-large |   3072    |   51.30   |    44.09    |   29.25   |   63.00    |   32.49   |   84.80    |   54.16   |
-|      baai-bge-m3       |   1024    |   56.56   |  **45.76**  |   40.25   |   46.00    |   35.54   | **94.09**  | **77.73** |
+# Tokenize input
+input_text = tokenizer('sample text', return_tensors='np')
 
-Notes: `*`, use the text-matching adapter
+# ONNX session
+model_path = 'jina-embeddings-v3/onnx/model.onnx'
+session = onnxruntime.InferenceSession(model_path)
 
+# Prepare inputs for ONNX model
+task_type = 'text-matching'
+task_id = np.array(config.lora_adaptations.index(task_type), dtype=np.int64)
+inputs = {
+    'input_ids': input_text['input_ids'],
+    'attention_mask': input_text['attention_mask'],
+    'task_id': task_id
+}
 
-#### Matryoshka Embeddings
+# Run model
+outputs = session.run(None, inputs)
+```
 
-|  Dimension  |  Retrieval  |   STS   |
-|:-----------:|:-----------:|:-------:|
-|     32      |    52.54    |  76.35  |
-|     64      |    58.54    |  77.03  |
-|     128     |    61.64    |  77.43  |
-|     256     |    62.72    |  77.56  |
-|     512     |    63.16    |  77.59  |
-|     768     |    63.3     |  77.59  |
-|    1024     |    63.35    |  77.58  |
+</p>
+</details>
 
-For a comprehensive evaluation and detailed metrics, please refer to the full paper available here (coming soon).
 
 ## Contact
 
@@ -21745,5 +21729,14 @@ Join our [Discord community](https://discord.jina.ai) and chat with other commun
 If you find `jina-embeddings-v3` useful in your research, please cite the following paper:
 
 ```bibtex
+@misc{sturua2024jinaembeddingsv3multilingualembeddingstask,
+      title={jina-embeddings-v3: Multilingual Embeddings With Task LoRA}, 
+      author={Saba Sturua and Isabelle Mohr and Mohammad Kalim Akram and Michael Günther and Bo Wang and Markus Krimmel and Feng Wang and Georgios Mastrapas and Andreas Koukounas and Andreas Koukounas and Nan Wang and Han Xiao},
+      year={2024},
+      eprint={2409.10173},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2409.10173}, 
+}
 
 ```

custom_st.py

@@ -91,19 +91,19 @@ class Transformer(nn.Module):
             self.auto_model.config.tokenizer_class = self.tokenizer.__class__.__name__
 
     def forward(
-        self, features: Dict[str, torch.Tensor], task_type: Optional[str] = None
+        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
     ) -> Dict[str, torch.Tensor]:
         """Returns token_embeddings, cls_token"""
-        if task_type and task_type not in self._lora_adaptations:
+        if task and task not in self._lora_adaptations:
             raise ValueError(
-                f"Unsupported task '{task_type}'. "
+                f"Unsupported task '{task}'. "
                 f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
-                f"Alternatively, don't pass the `task_type` argument to disable LoRA."
+                f"Alternatively, don't pass the `task` argument to disable LoRA."
             )
 
         adapter_mask = None
-        if task_type:
-            task_id = self._adaptation_map[task_type]
+        if task:
+            task_id = self._adaptation_map[task]
             num_examples = features['input_ids'].size(0)
             adapter_mask = torch.full(
                 (num_examples,), task_id, dtype=torch.int32, device=features['input_ids'].device

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17ca06efd886a065d0081912b04c9e27ef5086a9dd09659cce32aa9c84587f23
+size 1144685320

modules.json

@@ -4,7 +4,7 @@
     "name": "0",
     "path": "",
     "type": "custom_st.Transformer",
-    "kwargs": ["task_type"]
+    "kwargs": ["task"]
   },
   {
     "idx": 1,

onnx/_roberta_encoder_layers.0_mixer_rotary_emb_Constant_6_attr__value

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20472209ecbc9a1cea36feaa7e507db402e879d4fc51ebaf5f3cafefc4a7c4a9
+size 1048832

onnx/_roberta_encoder_layers.0_mixer_rotary_emb_Constant_attr__value

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7249c502306a0d781150401d414af5736ec47bed9f96395e5a3b504159342ab1
+size 1048832

onnx/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8586a2e2981d6ccf9d26104d926b92ca45c8c33a55fd146c62519b59aeaf7d2c
+size 3113411

onnx/roberta.embeddings.word_embeddings.parametrizations.weight.0.lora_A

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e331476c75a983801498b3a97c6e8599359b5002127a6802f94a11f9b67698b3
+size 20000160

onnx/roberta.embeddings.word_embeddings.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:827e52a1bff60a42513a0c2ded3bb3db168c973d473ab1a4e3281ec2060b4d98
+size 1024008192

onnx/roberta.encoder.layers.0.mixer.Wqkv.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:062b6c3a1067b86638020bcebe33f19e3d0a94d86ca95091fb9f465bf76b5c69
+size 12582912

onnx/roberta.encoder.layers.0.mixer.out_proj.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a17c40b5eba6671395bcad41b488730162eb8f26aea126d2de8f76857053adb5
+size 4194304

onnx/roberta.encoder.layers.0.mlp.fc1.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:841cf3f67fbb99cad22eb25d469b1300e1246dc62ae7e651ef7a50c37386098e
+size 16777216

onnx/roberta.encoder.layers.0.mlp.fc2.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a27c74c0d67d8a96ea66a48ca382d4a0a3ebf5354b1ce0bd75a1c8ef2f815ee
+size 16777216

onnx/roberta.encoder.layers.1.mixer.Wqkv.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1dbe254f0ec39a45b80c711a392141e8d4ff1d54e1aafb0c917855cf6b46612
+size 12582912

onnx/roberta.encoder.layers.1.mixer.out_proj.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1af212a115f6af2f0cd3d6430c71d53ea022cda186ed9097fcda06d243471598
+size 4194304

onnx/roberta.encoder.layers.1.mlp.fc1.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:083cd922bbe4bb93359461a08ca3ab1c80feb9ba08ac9332f07a2f204549fe0f
+size 16777216

onnx/roberta.encoder.layers.1.mlp.fc2.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:004cdb4408206f4a5c0d9c2e0a829ec0096c5ad98e9bbe93af8c0104af0d9e4a
+size 16777216

onnx/roberta.encoder.layers.10.mixer.Wqkv.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a0013b2bf24bdd46723a059c3bde4eb7347d9d1f894459ba42050209e6ecc7d
+size 12582912

onnx/roberta.encoder.layers.10.mixer.out_proj.parametrizations.weight.original

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3bb88e776197f7986730c10af22bbaac3de27496471e5e1a93ba46fe6f3e1d94
+size 4194304