File size: 1,030 Bytes
6491233 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
#!pip install uv
#!uv pip install transformers torch onnx
# Импорты
from transformers import AutoTokenizer, AutoModel
import torch
import onnx
tokenizer = AutoTokenizer.from_pretrained("ai-forever/sbert_large_nlu_ru")
model = AutoModel.from_pretrained("ai-forever/sbert_large_nlu_ru")
# Экспорт в ONNX
dummy_input = tokenizer("Тест", return_tensors="pt", padding=True, truncation=True)
torch.onnx.export(model,
(dummy_input['input_ids'], dummy_input['attention_mask']),
"sbert_large_nlu_ru.onnx",
input_names=["input_ids", "attention_mask"],
output_names=["last_hidden_state"],
dynamic_axes={
"input_ids": {0: "batch", 1: "sequence"},
"attention_mask": {0: "batch", 1: "sequence"},
"last_hidden_state": {0: "batch", 1: "sequence", 2: "features"}
})
# from google.colab import files
# files.download("/content/sbert_large_nlu_ru.onnx")
|