Update app.py

app.py

@@ -1,17 +1,28 @@
 import gradio as gr
 from transformers import pipeline
 
-model_name = "ieuniversity/flirty_classifier"
-classifier = pipeline("text-classification", model=model_name)
+classifier_model_name = "ieuniversity/flirty_classifier"
+paraphraser_model_name = "ieuniversity/flirty_paraphraser"
 
-def classify_text(text):
-    output = classifier(text)
-    label = output[0]["label"]
-    score = output[0]["score"]
-    return label
+classifier = pipeline("text-classification", model=classifier_model_name)
+paraphraser = pipeline("text2text-generation", model=paraphraser_model_name)
 
-gr.Interface(fn=classify_text, 
+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",
-             description="Enter some text and get a prediction for whether it's flirty or not.").launch()
+             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()