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 👍 Thanks for the awesome library!

[SidJain1412]

Created a PR regarding this: #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。