Generate torchscript for Sentence-Bert with sentence inputs

Hi There,

I’m a newbie to creating torch scripts from models. I have found this helpful page that shows how to export a PyTorch model for ‘BertModel’ for inputs at the token level using JIT & TRACE.

However, I want to export the SBert model which is also PyTorch based. But the inputs that SBert gets are sentences.

Below is a sample code I have put together, my problem is that I do not know how should I create the example_inputs here. Shall I convert sentences directly to tensor? or tokenize them first? I appreciate your input on this. Thanks

import torch
from sentence_transformers import SentenceTransformer

sbert_model        = SentenceTransformer('paraphrase-mpnet-base-v2')
sbert_model.eval()

# Creating the trace
traced_model = torch.jit.trace(sbert_model, example_inputs = [])
torch.jit.save(traced_model, "traced_bsert.pt")