File size: 2,827 Bytes
5120311 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
from sentence_transformers import SentenceTransformer
import torch
from transformers import AutoTokenizer
def convert_onnx():
# model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
model = SentenceTransformer('model/distiluse-base-multilingual-cased-v2')
torch.save(model, "model/distiluse-base-multilingual-cased-v2.pt")
model = torch.load("model/distiluse-base-multilingual-cased-v2.pt")
tokenizer = AutoTokenizer.from_pretrained('model/distiluse-base-multilingual-cased-v2')
# tokenizer = AutoTokenizer.from_pretrained('deepset/roberta-base-squad2')
# lst_input = [" Experts predict China will soon attack Taiwan"]
lst_input = ["Pham Minh Chinh is Vietnam's Prime Minister"]
x = tokenizer(lst_input, padding="max_length", truncation=True)
print(x)
save_path = 'tensorRT/models/distiluse-base-multilingual-cased-v2.onnx'
torch.onnx.export(model, (torch.tensor(x['input_ids'], dtype=torch.long),torch.tensor(x['attention_mask'], dtype=torch.long)), save_path, export_params=True, opset_version=13, do_constant_folding=True,
input_names = ['input_ids', 'attention_mask'],
output_names = ['output'],
dynamic_axes={'input_ids' : {0 : 'batch_size'}, 'attention_mask': {0 : 'batch_size'},
'output' : {0 : 'batch_size'}}
)
def convert_onnx_(model_name= "model/model-sup-simcse-vn", pt_model = 'model-sup-simcse-vn', max_length = 256, save_path = "tensorRT/models/tensorRT/models/model-sup-simcse-vn.onnx"):
model = SentenceTransformer(model_name)
torch.save(model, f"model/{pt_model}.pt")
model = torch.load(f"model/{pt_model}.pt")
tokenizer = AutoTokenizer.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained('deepset/roberta-base-squad2')
# lst_input = [" Experts predict China will soon attack Taiwan"]
lst_input = ["Pham Minh Chinh is Vietnam's Prime Minister"]
x = tokenizer(lst_input, padding="max_length", truncation=True, max_length=256)
print(x)
torch.onnx.export(model, (torch.tensor(x['input_ids'], dtype=torch.long),torch.tensor(x['attention_mask'], dtype=torch.long), torch.tensor(x['token_type_ids'], dtype=torch.long)), save_path, export_params=True, opset_version=13, do_constant_folding=True,
input_names = ['input_ids', 'attention_mask','token_type_ids'],
output_names = ['output'],
dynamic_axes={'input_ids' : {0 : 'batch_size'}, 'attention_mask': {0 : 'batch_size'},
'output' : {0 : 'batch_size'}}
)
if __name__ == '__main__':
convert_onnx_() |