import gradio as gr
import torch
from transformers import PegasusForConditionalGeneration, PegasusTokenizer
from sentence_splitter import SentenceSplitter, split_text_into_sentences

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=10000,
                                            return_tensors="pt").to(torch_device)
    translated = model.generate(**batch, num_beams=10, num_return_sequences=num_return_sequences,
                                temperature=1.5)
    tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
    return tgt_text


def get_response_from_text(
        context="I am a student at the University of Washington. I am taking a course called Data Science."):
    splitter = SentenceSplitter(language='en')
    sentence_list = splitter.split(context)

    paraphrase = []

    for i in sentence_list:
        a = get_response(i, 1)
        paraphrase.append(a)
    paraphrase2 = [' '.join(x) for x in paraphrase]
    paraphrase3 = [' '.join(x for x in paraphrase2)]
    paraphrased_text = str(paraphrase3).strip('[]').strip("'")
    return paraphrased_text

def greet(context):
    return get_response_from_text(context)

examples = [["Begin your professional career by learning data science skills with Data science Dojo, a globally recognized e-learning platform where we teach students how to learn data science, data analytics, machine learning and more."], ["Paraphrasing is a way to express the meaning of something using different words in order to make it easier for audience to understand. This app uses natural language processing to create paraphrasing of your input text."]]

css = """
footer {display:none !important}
.output-markdown{display:none !important}

.gr-button-primary {
    z-index: 14;
    height: 30px;
    width: 113px;
    left: 0px;
    top: 0px;
    padding: 0px;
    cursor: pointer !important; 
    background: none rgb(17, 20, 45) !important;
    border: none !important;
    text-align: center !important;
    font-family: Poppins !important;
    font-size: 14px !important;
    font-weight: 500 !important;
    color: rgb(255, 255, 255) !important;
    line-height: 1 !important;
    border-radius: 12px !important;
    transition: box-shadow 200ms ease 0s, background 200ms ease 0s !important;
    box-shadow: none !important;
}
.gr-button-primary:hover{
    z-index: 14;
    height: 30px;
    width: 113px;
    left: 0px;
    top: 0px;
    padding: 0px;
    cursor: pointer !important;
    background: none rgb(37, 56, 133) !important;
    border: none !important;
    text-align: center !important;
    font-family: Poppins !important;
    font-size: 14px !important;
    font-weight: 500 !important;
    color: rgb(255, 255, 255) !important;
    line-height: 1 !important;
    border-radius: 12px !important;
    transition: box-shadow 200ms ease 0s, background 200ms ease 0s !important;
    box-shadow: rgb(0 0 0 / 23%) 0px 1px 7px 0px !important;
}
.hover\:bg-orange-50:hover {
    --tw-bg-opacity: 1 !important;
    background-color: rgb(229,225,255) !important;
}
"""

demo = gr.Interface(fn=greet, inputs=gr.Textbox(lines=3, placeholder="Enter sample text here", label="Original text"), outputs=gr.Textbox(label="Paraphrasing"), title="Paraphrasing | Data Science Dojo", examples=examples, css=css)

demo.launch()