jupyterjazz
commited on
Commit
•
4f519a4
1
Parent(s):
712be5f
Update README.md
Browse files
README.md
CHANGED
@@ -118,7 +118,42 @@ We used the `unique()` function to identify unique task types in a batch, which
|
|
118 |
|
119 |
|
120 |
|
|
|
121 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
122 |
|
123 |
|
124 |
|
|
|
118 |
|
119 |
|
120 |
|
121 |
+
## Code
|
122 |
|
123 |
+
```python
|
124 |
+
import torch
|
125 |
+
from transformers import AutoModel, AutoTokenizer
|
126 |
+
import torch.onnx
|
127 |
+
|
128 |
+
|
129 |
+
model = AutoModel.from_pretrained('/home/admin/saba/jina-embeddings-v3', trust_remote_code=True, use_flash_attn=False)
|
130 |
+
model.eval()
|
131 |
+
|
132 |
+
onnx_path = "/home/admin/saba/jina-embeddings-v3/onnx/model.onnx"
|
133 |
+
|
134 |
+
tokenizer = AutoTokenizer.from_pretrained('/home/admin/saba/jina-embeddings-v3')
|
135 |
+
inputs = tokenizer(["jina", 'ai'], return_tensors="pt", padding='longest')
|
136 |
+
inps = inputs['input_ids']
|
137 |
+
mask = inputs['attention_mask']
|
138 |
+
task_id = 2
|
139 |
+
|
140 |
+
|
141 |
+
torch.onnx.export(
|
142 |
+
model,
|
143 |
+
(inps, mask, task_id),
|
144 |
+
onnx_path,
|
145 |
+
export_params=True,
|
146 |
+
do_constant_folding=True,
|
147 |
+
input_names = ['input_ids', 'attention_mask', 'task_id'],
|
148 |
+
output_names = ['text_embeds'],
|
149 |
+
opset_version=16,
|
150 |
+
dynamic_axes={
|
151 |
+
'input_ids' : {0 : 'batch_size', 1: 'sequence_length'},
|
152 |
+
'attention_mask' : {0 : 'batch_size', 1: 'sequence_length'},
|
153 |
+
'text_embeds' : {0 : 'batch_size'}
|
154 |
+
},
|
155 |
+
)
|
156 |
+
```
|
157 |
|
158 |
|
159 |
|