Update app.py

app.py

@@ -1,37 +1,52 @@
 import gradio as gr
 from transformers import pipeline
 import re
+from keybert import KeyBERT
+import torch
 
 text_pipe = pipeline("text-classification", model="karalif/myTestModel", return_all_scores=True)
 
+def get_prediction(text):
+    # Tokenize the input text
+    encoding = new_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=200)
+    encoding = {k: v.to(new_model.device) for k, v in encoding.items()}
+
+    with torch.no_grad():
+        outputs = new_model(**encoding)
+
+    logits = outputs.logits
+    sigmoid = torch.nn.Sigmoid()
+    probs = sigmoid(logits.squeeze().cpu()).numpy()
+
+    # Initialize KeyBERT
+    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)
+
+    response = "MY PREDICTION:<br>"
+    labels = ['Politeness', 'Toxicity', 'Sentiment', 'Formality']
+
+    for i, label in enumerate(labels):
+        response += f"{label}: {probs[i]*100:.1f}%<br>"
+    
+    response += "<br>INFLUENTIAL KEYWORDS:<br>"
+    for keyword, score in keywords:
+        response += f"{keyword} (Score: {score:.2f})<br>"
+
+    return response
+
 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)
     greeting_feedback = ""
     
-    results = text_pipe(text)
-    all_scores = results[0]
-    response = ""
-    for result in all_scores:
-        label = result['label']
-        score = result['score']
-        if label == "Politeness":
-            response += f"<span style='background-color:#b8e994; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Green
-        elif label == "Sentiment":
-            response += f"<span style='background-color:#fff3cd; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Yellow
-        elif label == "Formality":
-            response += f"<span style='background-color:#bee5eb; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Blue
-        elif label == "Toxicity":
-            response += f"<span style='background-color:#f8d7da; color:black;'>{label}</span>: {score:.3f}<br>"  # Light Red
-    
     if not re.match(greeting_pattern, text, re.IGNORECASE):
         greeting_feedback = "Heilsaðu dóninn þinn<br>"
 
-    response += greeting_feedback
+    response = prediction_output + "<br>" + greeting_feedback
     
     return response
 
-
 description_html = """
     <center>
         <img src='http://www.ru.is/media/HR_logo_vinstri_transparent.png' width='250' height='auto'>