Update app.py

app.py

@@ -2,13 +2,17 @@ import gradio as gr
 import re
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
-from keybert import KeyBERT
+import shap
+import numpy as np
 
 # Initialize your model and tokenizer here
 model_identifier = "karalif/myTestModel"
 new_model = AutoModelForSequenceClassification.from_pretrained(model_identifier)
 new_tokenizer = AutoTokenizer.from_pretrained(model_identifier)
 
+# SHAP Explainer Initialization
+explainer = shap.Explainer(new_model, new_tokenizer)
+
 def get_prediction(text):
     # Tokenize the input text
     encoding = new_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=200)
@@ -21,9 +25,14 @@ def get_prediction(text):
     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)
+    # Generate SHAP values
+    shap_values = explainer([text])
+
+    # Extracting top SHAP values and their corresponding tokens
+    top_shap_values = np.abs(shap_values.values).mean(0).sum(-1)
+    top_tokens_indices = np.argsort(-top_shap_values)[:5]  # Getting indices of top 5 tokens
+    top_tokens = [new_tokenizer.convert_ids_to_tokens(encoding['input_ids'][0][idx].item()) for idx in top_tokens_indices]
+    top_shap_scores = top_shap_values[top_tokens_indices]
 
     # Prepare the HTML output with labels and their probabilities
     response = ""
@@ -34,10 +43,10 @@ def get_prediction(text):
         response += f"<span style='background-color:{colors[i]}; color:black;'>{label}</span>: {probs[i]*100:.1f}%<br>"
     
     influential_keywords = "INFLUENTIAL KEYWORDS:<br>"
-    for keyword, score in keywords:
-        influential_keywords += f"{keyword} (Score: {score:.2f})<br>"
+    for token, score in zip(top_tokens, top_shap_scores):
+        influential_keywords += f"{token} (Score: {score:.2f})<br>"
     
-    return response, keywords, influential_keywords
+    return response, list(zip(top_tokens, top_shap_scores)), 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)"
@@ -48,7 +57,7 @@ def predict(text):
     # 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>")
+        modified_input = re.sub(rf"(\b{keyword}\b)", r"<span style='color:green;'>\1</span>", modified_input, flags=re.IGNORECASE)
     
     if not re.match(greeting_pattern, text, re.IGNORECASE):
         greeting_feedback = "OTHER FEEDBACK:<br>Heilsaðu dóninn þinn<br>"