Hi,
We have finetuned distill-pegasus-cnn-16-4 5 summarization model on our own data and results look good. However, when we want to deploy it for a real-time production use case - it is taking huge time on ml.c5.xlarge CPU (around 13seconds per document in a sequence). We tried a g4dn.xlarge 4 GPU for inference and it is taking around 1.7seconds for one document in a sequence. Inference on a GPU is a costly affair. Could anyone suggest me if anyone have any ideas to make faster with less cost.
I have tried “onnx runtime”, TorchScript but both of them dont have support for pegasus model. Do we have any timelines for supporting pegasus. I tried num_beams as well - above numbers are after using them.
Inference code:
import torch
from transformers import PegasusForConditionalGeneration, PegasusTokenizer
src_text = [
“”" PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."""
]
model_name = “sshleifer/distill-pegasus-cnn-16-4”
torch_device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
tokenizer = PegasusTokenizer.from_pretrained(model_name)
model = PegasusForConditionalGeneration.from_pretrained(model_name).to(torch_device)
inputs = tokenizer.batch_encode_plus(src_text, truncation=True, padding=“longest”, return_tensors=“pt”).to(torch_device)
translated = model.generate(inputs[“input_ids”], num_beams=2)
output = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in translated][0]
print(output)
Thanks in advance for your help. Apologies if I am missing something obvious.
Regards,
Karthik