label map

app.py

@@ -4,6 +4,8 @@ from gradio.components import Textbox
 
 # Load the sentiment analysis pipeline with DistilBERT
 distilbert_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
+label_map = {"POSITIVE":"OTHER", "NEGATIVE":"SENSITIVE"}
+
 
 def predict_sentiment(text):
     """
@@ -12,15 +14,15 @@ def predict_sentiment(text):
     :return: str, predicted sentiment and confidence score.
     """
     result = distilbert_pipeline(text)[0]
-    label = result['label']
+    label = label_map[result['label']]
     score = result['score']
-    return f"Sentiment: {label}, Confidence: {score:.2f}"
+    return f"TAG: {label}, Confidence: {score:.2f}"
 
 input1 = Textbox(lines=2, placeholder="Type your text here...")
 
 # Create a Gradio interface
 iface = gr.Interface(fn=predict_sentiment,
-                     inputs=input1
+                     inputs=input1,
                      outputs="text",
                      title="Talk2Loop Sensitive statement tags",
                      description="This model predicts the sensitivity of the input text. Enter a sentence to see if it's sensitive or not.")