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()