Update README.md

README.md

@@ -118,7 +118,42 @@ We used the `unique()` function to identify unique task types in a batch, which
 
 
 
+## Code
 
+```python
+import torch
+from transformers import AutoModel, AutoTokenizer
+import torch.onnx
+
+
+model = AutoModel.from_pretrained('/home/admin/saba/jina-embeddings-v3', trust_remote_code=True, use_flash_attn=False)
+model.eval()
+
+onnx_path =  "/home/admin/saba/jina-embeddings-v3/onnx/model.onnx"
+
+tokenizer = AutoTokenizer.from_pretrained('/home/admin/saba/jina-embeddings-v3')
+inputs = tokenizer(["jina", 'ai'], return_tensors="pt", padding='longest')
+inps = inputs['input_ids']
+mask = inputs['attention_mask']
+task_id = 2
+
+
+torch.onnx.export(
+    model,
+    (inps, mask, task_id),
+    onnx_path,
+    export_params=True,
+    do_constant_folding=True,
+    input_names = ['input_ids', 'attention_mask', 'task_id'],
+    output_names = ['text_embeds'],
+    opset_version=16,
+    dynamic_axes={
+        'input_ids' : {0 : 'batch_size', 1: 'sequence_length'},
+        'attention_mask' : {0 : 'batch_size', 1: 'sequence_length'},
+        'text_embeds' : {0 : 'batch_size'}
+    },
+)
+```