Update app.py

app.py

@@ -2,17 +2,13 @@ import gradio as gr
 import re
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
-import shap
-import numpy as np
+from keybert import KeyBERT
 
 # 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)
@@ -25,14 +21,9 @@ def get_prediction(text):
     sigmoid = torch.nn.Sigmoid()
     probs = sigmoid(logits.squeeze().cpu()).numpy()
 
-    # 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]
+    # 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)
 
     # Prepare the HTML output with labels and their probabilities
     response = ""
@@ -43,10 +34,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 token, score in zip(top_tokens, top_shap_scores):
-        influential_keywords += f"{token} (Score: {score:.2f})<br>"
+    for keyword, score in keywords:
+        influential_keywords += f"{keyword} (Score: {score:.2f})<br>"
     
-    return response, list(zip(top_tokens, top_shap_scores)), influential_keywords
+    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)"
@@ -57,7 +48,7 @@ def predict(text):
     # Highlight the keywords in the input text
     modified_input = text
     for keyword, _ in keywords:
-        modified_input = re.sub(rf"(\b{keyword}\b)", r"<span style='color:green;'>\1</span>", modified_input, flags=re.IGNORECASE)
+        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>"