Sentence-transformers: sentence bert model in onnx format

[rachel2011]

I would like to convert sentence bert model from pytorch to tensorflow use onnx, and tried to follow the standard onnx procedure for converting a pytorch model. But I'm having difficulty determining the onnx input arguments for sentence bert model, I encounter TypeError: forward() takes 2 positional arguments but 4 were given. Suggestions appreciated!
model = SentenceTransformer('output/continue_training_model')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
dummy_input0 = torch.LongTensor(batch_size, max_seq_length).to(device)
dummy_input1 = torch.LongTensor(batch_size, max_seq_length).to(device)
dummy_input2 = torch.LongTensor(batch_size, max_seq_length).to(device)
torch.onnx.export(model,(dummy_input0, dummy_input1,dummy_input2), onnx_file_name, verbose=True)

[nreimers]

Sadly I never worked with onnx.

In SentenceTransformer, the forward function takes in one argument: features (and the second in python is self).

Features is a dictionary, that contains the different features, for example, token ids, word weights, attention values, token_type_ids.

For the BERT model, I think your input must look like this:
input_features = {'input_ids': dummy_input0, 'token_type_ids': dummy_input1, 'input_mask': dummy_input2}

And then:

torch.onnx.export(model,input_features, onnx_file_name, verbose=True)

[ycgui]

@rachel2011 Did you get any solutions to successfully convert sentence bert model to onnx format?

I'm also wondering how to feed text inputs into a converted onnx model. Could we do something similar as

sentences = ['This framework generates embeddings for each input sentence',
    'Sentences are passed as a list of string.', 
    'The quick brown fox jumps over the lazy dog.']
sentence_embeddings = model.encode(sentences)

by replacing the model to the converted onnx model? Any ideas?

[nreimers]

Hi @ycgui
I started to add the models to HuggingFace Models Hub:
https://huggingface.co/sentence-transformers

Huggingface also provides methods / scripts to convert models to ONNX.

I hope this helps.

[ycgui]

Thanks @nreimers. This is awesome!

[SidJain1412]

Hey @ycgui I would be really thankful if you could share the code you used to convert the models to ONNX and then how you can encode sentences using that model.
Thank you in advance

[SidJain1412]

Hey @ycgui I would be really thankful if you could share the code you used to convert the models to ONNX and then how you can encode sentences using that model.
Thank you in advance

As per my understanding shouldn't the pooled output after running the ONNX model match the output of encode using SentenceTransformers?
This doesn't seem to be the case in my testing. (using the same model and tokenizer in both cases)

@nreimers , I would appreciate your help greatly

[nreimers]

Sadly I am not aware of ONNX format.

Here you can see an example how to load the models with native transformers code and how to apply mean pooling correctly (watch out for padding tokens):
https://huggingface.co/sentence-transformers/bert-base-nli-mean-tokens

[SidJain1412]

Thank you a lot @nreimers , applying the mean pooling correctly I was able to get sentence embeddings as expected!

I am thinking about making a small tutorial Notebook on how to use sentence transformers with ONNX + benchmarking it against the traditional PyTorch model, would that be useful in Examples?

[nreimers]

Yes, it would be great to have such tutorial.

[SidJain1412]

Awesome, I'll make a PR regarding that soon :+1: Thanks for the awesome library!

[SidJain1412]

Created a PR regarding this: https://github.com/UKPLab/sentence-transformers/pull/386

[nreimers]

great, I will have a look

[cantwbr]

@rachel2011: My response might be a bit late. I think the keys in your dictionary are wrong. For sentence-transformers in version 0.3.7.2, I downloaded the models (like _bert-base-nli-mean-tokens_) from here. Then I used
input_features = {'input_ids': input_ids, 'token_type_ids': input_type_ids, 'attention_mask': input_mask}
and
torch.onnx.export(model,input_features, onnx_file_name, verbose=True)
to export the sentence bert model.

[SidJain1412]

@cantwbr A simpler way is to use from transformers.convert_graph_to_onnx import convert that can be used to convert to an ONNX model. Refer to this PR.

[codingliuyg]

@rachel2011: My response might be a bit late. I think the keys in your dictionary are wrong. For sentence-transformers in version 0.3.7.2, I downloaded the models (like _bert-base-nli-mean-tokens_) from here. Then I used
input_features = {'input_ids': input_ids, 'token_type_ids': input_type_ids, 'attention_mask': input_mask}
and
torch.onnx.export(model,input_features, onnx_file_name, verbose=True)
to export the sentence bert model.

hello,how do you define the input_ids、input_type_ids and input_mask，can you show your demo? My code as shown below does not pass。

batch_size=1
max_seq_length=128
device = torch.device("cuda")
model.to(device)
dummy_input0 = torch.LongTensor(batch_size, max_seq_length).to(device)
dummy_input1 = torch.LongTensor(batch_size, max_seq_length).to(device)
dummy_input2 = torch.LongTensor(batch_size, max_seq_length).to(device)
input_features = {'input_ids': dummy_input0, 'token_type_ids': dummy_input1, 'attention_mask': dummy_input2}

@cantwbr Looking forward to your reply，thanks。

[cantwbr]

@codingliuyg: I didn't use LongTensors. Instead I generated ones-tensors.

input_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_type_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_mask = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_features = {'input_ids': input_ids, 'token_type_ids': input_type_ids, 'attention_mask': input_mask}
torch.onnx.export(model,input_features, onnx_file_name, verbose=True)

[codingliuyg]

@codingliuyg: I didn't use LongTensors. Instead I generated ones-tensors.

input_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_type_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_mask = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_features = {'input_ids': input_ids, 'token_type_ids': input_type_ids, 'attention_mask': input_mask}
torch.onnx.export(model,input_features, onnx_file_name, verbose=True)

@cantwbr thank you for your replay.After changing my code as following，A error occurs。Do you know what happened ?Is there anything wrong with my code? thank you .

_model = SentenceTransformer('roberta-base-nli-stsb-mean-tokens',device='cpu')
batch_size=1
max_seq_length=128
device = torch.device("cpu")
model.to(device)
input_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_type_ids = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_mask = torch.ones(batch_size, max_seq_length, dtype=torch.long).to(device)
input_features = {'input_ids': input_ids, 'token_type_ids': input_type_ids, 'attention_mask': input_mask}
onnx_path = "onnx_model_name.onnx"
torch.onnx.export(model, input_features, onnx_path)_

error:
2020-12-17 21:24:17 - Load pretrained SentenceTransformer: roberta-base-nli-stsb-mean-tokens
2020-12-17 21:24:17 - Load SentenceTransformer from folder: roberta-base-nli-stsb-mean-tokens
Traceback (most recent call last):
File "embedding_reduce_cp.py", line 45, in
torch.onnx.export(model, input_features, onnx_path)
File "/opt/conda/lib/python3.6/site-packages/torch/onnx/__init__.py", line 230, in export
custom_opsets, enable_onnx_checker, use_external_data_format)
File "/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py", line 92, in export
use_external_data_format=use_external_data_format)
File "/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py", line 538, in _export
fixed_batch_size=fixed_batch_size)
File "/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py", line 374, in _model_to_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "/opt/conda/lib/python3.6/site-packages/torch/onnx/utils.py", line 327, in _trace_and_get_graph_from_model
torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/__init__.py", line 135, in _get_trace_graph
outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(args, kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 726, in _call_impl
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_trace.py", line 116, in forward
self._force_outplace,
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_trace.py", line 102, in wrapper
outs.append(self.inner(trace_inputs))
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 724, in _call_impl
result = self._slow_forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 708, in _slow_forward
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py", line 117, in forward
input = module(input)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 724, in _call_impl
result = self._slow_forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 708, in _slow_forward
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/sentence_transformers/models/Transformer.py", line 36, in forward
output_states = self.auto_model(features)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 724, in _call_impl
result = self._slow_forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 708, in _slow_forward
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/transformers/modeling_roberta.py", line 675, in forward
input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 724, in _call_impl
result = self._slow_forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 708, in _slow_forward
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/transformers/modeling_roberta.py", line 119, in forward
token_type_embeddings = self.token_type_embeddings(token_type_ids)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 724, in _call_impl
result = self._slow_forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 708, in _slow_forward
result = self.forward(input, *kwargs)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/sparse.py", line 126, in forward
self.norm_type, self.scale_grad_by_freq, self.sparse)
File "/opt/conda/lib/python3.6/site-packages/torch/nn/functional.py", line 1837, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
IndexError: index out of range in self

[cantwbr]

@codingliuyg: I tried exporting _roberta-base-nli-stsb-mean-tokens_ and got a similar error. I resolved it by changing torch.ones to torch.zeros.
I used torch 1.4.0, sentence-transformers 0.3.7.2 in Python 3.8.

[shar999]

@cantwbr can you share the code that you used for inference?

[cantwbr]

@shar999: For inference using the ONNX file, I reuse the featurizer of the original SentenceBert package from

from sentence_transformers.datasets import EncodeDataset

import onnxruntime as ort
from sentence_transformers import SentenceTransformer
from torch.utils.data import DataLoader

The inference:

model = SentenceTransformer(model_name_or_path)
fused_bert_session = ort.InferenceSession(onnx_file_name)

batch_size = 1
is_pretokenized = False
num_workers = 0
pad_to = 128

all_embeddings = []
length_sorted_idx = np.argsort([model._text_length(sen) for sen in sentences])
sentences_sorted = [sentences[idx] for idx in length_sorted_idx]
inp_dataset = EncodeDataset(sentences_sorted, model=model, is_tokenized=is_pretokenized)
inp_dataloader = DataLoader(inp_dataset, batch_size=batch_size, collate_fn=model.smart_batching_collate_text_only,
                                num_workers=num_workers, shuffle=False)

iterator = inp_dataloader
for features in iterator:
   for feature_name in features:
         pad_amt = features[feature_name].size(-1)
         if pad_amt != 0:
             features[feature_name] = nn.functional.pad(features[feature_name], (0, pad_to - pad_amt), value=0)

    e_input_ids = features["input_ids"].cpu().numpy()
    e_input_type_ids = features["token_type_ids"].cpu().numpy()
    e_input_mask = features["attention_mask"].cpu().numpy()

    results = fused_bert_session.run(None,{"input_ids": e_input_ids, "token_type_ids": e_input_type_ids, "attention_mask": e_input_mask,})
    # results[5] contains the sentence_embeddings
    all_embeddings.append(results[5])

The embeddings are located in all_embeddings.