import gradio as gr
from transformers import pipeline

classifier_model_name = "ieuniversity/flirty_classifier"
paraphraser_model_name = "ieuniversity/flirty_paraphraser"

classifier = pipeline("text-classification", model=classifier_model_name)
paraphraser = pipeline("text2text-generation", model=paraphraser_model_name)

def classify_and_paraphrase_text(text):
    classification_output = classifier(text)
    label = classification_output[0]["label"]
    score = classification_output[0]["score"]
    
    if label == "flirty":
        paraphrase_output = paraphraser(text, max_length=100, do_sample=True, temperature=0.8)
        paraphrase_text = paraphrase_output[0]["generated_text"]
        return f"This message seems flirty. Here's another suggestion: '{paraphrase_text}'"
    else:
        paraphrase_output = paraphraser(text, max_length=100, do_sample=True, temperature=0.8)
        paraphrase_text = paraphrase_output[0]["generated_text"]
        return f"This message doesn't seem flirty. How about: '{paraphrase_text}'?"

gr.Interface(fn=classify_and_paraphrase_text, 
             inputs="text", 
             outputs="text",
             title="Flirty Classifier and Paraphraser",
             description="Enter some text and get a prediction for whether it's flirty or not. If it's not flirty, get a paraphrased suggestion.").launch()