Update app.py

app.py

@@ -25,33 +25,35 @@ 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>"
+    # Prepare the HTML output with labels and their probabilities
+    response = ""
     labels = ['Politeness', 'Toxicity', 'Sentiment', 'Formality']
     colors = ['#b8e994', '#f8d7da', '#fff3cd', '#bee5eb']  # Corresponding colors for labels
 
     for i, label in enumerate(labels):
         response += f"<span style='background-color:{colors[i]}; color:black;'>{label}</span>: {probs[i]*100:.1f}%<br>"
     
-    response += "<br>INFLUENTIAL KEYWORDS:<br>"
+    influential_keywords = "<br>INFLUENTIAL KEYWORDS:<br>"
     for keyword, score in keywords:
-        response += f"{keyword} (Score: {score:.2f})<br>"
-
-    return response
+        influential_keywords += f"{keyword} (Score: {score:.2f})<br>"
+    
+    return response, keywords, influential_keywords
 
 def predict(text):
     greeting_pattern = r"^(Halló|Hæ|Sæl|Góðan dag|Kær kveðja|Daginn|Kvöldið|Ágætis|Elsku)"
     
-    prediction_output = get_prediction(text)
+    prediction_output, keywords, influential_keywords = get_prediction(text)
     greeting_feedback = ""
     
+    # Highlight the keywords in the input text
+    modified_input = text
+    for keyword, _ in keywords:
+        modified_input = modified_input.replace(keyword, f"<span style='color:green;'>{keyword}</span>")
+    
     if not re.match(greeting_pattern, text, re.IGNORECASE):
         greeting_feedback = "OTHER FEEDBACK:<br>Heilsaðu dóninn þinn<br>"
 
-    # Prepend the input text to the response
-    input_text_display = f"INPUT:<br>{text}<br><br>"
-
-    response = input_text_display + prediction_output + "<br>" + greeting_feedback
+    response = f"INPUT:<br>{modified_input}<br><br>MY PREDICTION:<br>{prediction_output}<br>{influential_keywords}<br>{greeting_feedback}"
     
     return response