Spaces:
Running
Running
import gradio as gr | |
import torch | |
from transformers import PegasusForConditionalGeneration, PegasusTokenizer | |
model_name = 'tuner007/pegasus_paraphrase' | |
# torch_device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
tokenizer = PegasusTokenizer.from_pretrained(model_name) | |
model = PegasusForConditionalGeneration.from_pretrained(model_name) | |
def paraphrase(text): | |
from sentence_splitter import SentenceSplitter, split_text_into_sentences | |
import torch | |
from transformers import PegasusForConditionalGeneration, PegasusTokenizer | |
model_name = 'tuner007/pegasus_paraphrase' | |
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) | |
def get_response(input_text,num_return_sequences): | |
batch = tokenizer.prepare_seq2seq_batch([input_text],truncation=True,padding='longest',max_length=60, return_tensors="pt").to(torch_device) | |
translated = model.generate(**batch,max_length=60,num_beams=10, num_return_sequences=num_return_sequences, temperature=2.0) | |
tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True) | |
return tgt_text | |
splitter = SentenceSplitter(language='en') | |
sentence_list = splitter.split(text) | |
res = '' | |
for i in sentence_list: | |
a = get_response(i,1) | |
cur = '' | |
for j in a: | |
cur += j | |
cur += ' ' | |
cur += '.' | |
res += cur | |
return res | |
iface = gr.Interface(fn=paraphrase, inputs="text", outputs="text") | |
iface.launch() |