import gradio as gr
from parrot import Parrot
import torch

# Initialize the Parrot model
parrot = Parrot(model_tag="prithivida/parrot_paraphraser_on_T5")

# Define paraphrasing function
def paraphrase_text(input_text):
    if not input_text:
        return ["Please enter text to paraphrase."]
    
    # Generate paraphrases
    paraphrases = parrot.augment(input_phrase=input_text, diversity_ranker="levenshtein", do_diverse=True)
    
    if paraphrases:
        return [p for p in paraphrases]
    else:
        return ["No paraphrase generated. Try another input."]

# Create UI with Gradio
interface = gr.Interface(
    fn=paraphrase_text,
    inputs=gr.Textbox(lines=3, placeholder="Enter text to paraphrase"),
    outputs=gr.Textbox(lines=6),
    title="Parrot Paraphraser",
    description="An AI-powered paraphraser using the Parrot model. Enter a sentence, and it will generate multiple reworded versions."
)

if __name__ == "__main__":
    interface.launch()