Update app.py

app.py

@@ -5,7 +5,6 @@ import torch
 from keybert import KeyBERT
 
 # Initialize your model and tokenizer here
-# Replace 'your_model_path' and 'your_tokenizer_path' with the actual paths or model identifiers
 model_identifier = "karalif/myTestModel"
 new_model = AutoModelForSequenceClassification.from_pretrained(model_identifier)
 new_tokenizer = AutoTokenizer.from_pretrained(model_identifier)
@@ -26,11 +25,13 @@ def get_prediction(text):
     kw_model = KeyBERT()
     keywords = kw_model.extract_keywords(text, keyphrase_ngram_range=(1, 1), stop_words='english', use_maxsum=True, nr_candidates=20, top_n=5)
 
+    # Modify this part to include HTML span with background color
     response = "MY PREDICTION:<br>"
     labels = ['Politeness', 'Toxicity', 'Sentiment', 'Formality']
+    colors = ['#b8e994', '#f8d7da', '#fff3cd', '#bee5eb']  # Corresponding colors for labels
 
     for i, label in enumerate(labels):
-        response += f"{label}: {probs[i]*100:.1f}%<br>"
+        response += f"<span style='background-color:{colors[i]}; color:black;'>{label}</span>: {probs[i]*100:.1f}%<br>"
     
     response += "<br>INFLUENTIAL KEYWORDS:<br>"
     for keyword, score in keywords:
@@ -72,4 +73,4 @@ demo = gr.Interface(
     theme=gr.themes.Default(primary_hue="red", secondary_hue="pink")
 )
 
-demo.launch()
+demo.launch()