Updated rules to ignore punctuation

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import numpy as np
 import torch
+import re
 from nltk.stem import PorterStemmer
 from collections import defaultdict
 from transformers import pipeline
@@ -32,7 +33,8 @@ tab_two_examples = [[ex_str1, ex_key1],
 #                     ['The _ plane is composed of a two-dimensional hexagonal lattice of carbon atoms.']
 #                     ]
 
-ignore = ['a', 'an', 'the', 'is', 'and', 'or']
+ignore_str = ['a', 'an', 'the', 'is', 'and', 'or', '!', '(', ')', '-', '[', ']', '{', '}', ';', ':', "'", '"', '\\',
+              ',', '<', '>', '.', '/', '?', '@', '#', '$', '%', '^', '&', '*', '_', '~']
 
 
 def add_mask(text, lower_bound=0, index=None):
@@ -49,7 +51,7 @@ def add_mask(text, lower_bound=0, index=None):
     idx = np.random.randint(low=lower_bound, high=len(split_text), size=1).astype(int)[0]
     # Don't mask certain words
     num_iters = 0
-    while split_text[idx].lower() in ignore:
+    while split_text[idx].lower() in ignore_str:
         num_iters += 1
         idx = np.random.randint(len(split_text), size=1).astype(int)[0]
         if num_iters > 10:
@@ -220,8 +222,9 @@ def extract_keywords(text, queries):
         # Iterate through text and mask each token
         ps = PorterStemmer()
         top_scores = defaultdict(list)
-        top_k_range = 10
-        indices = [i for i, t in enumerate(text.split()) if t.lower() == query.lower()]
+        top_k_range = 30
+        text_no_punc = re.sub(r'[^\w\s]', '', text)
+        indices = [i for i, t in enumerate(text_no_punc.split()) if t.lower() == query.lower()]
         for i in indices:
             masked_text, masked = add_mask(text, index=i)
             res = scrambler(masked_text, temp=temp, top_k=top_k_range)
@@ -229,12 +232,14 @@ def extract_keywords(text, queries):
             sorted_keys = sorted(out, key=out.get)
             # If the key does not appear, floor its rank for that round
             for rank, token_str in enumerate(sorted_keys):
+                if token_str in ignore_str:
+                    continue
                 stemmed = ps.stem(token_str)
-                if token_str not in top_scores.keys():
-                    top_scores[stemmed].append(0)
                 norm_rank = rank / top_k_range
                 top_scores[stemmed].append(norm_rank)
-
+            for key in top_scores.keys():
+                if key not in out.keys():
+                    top_scores[key].append(0)
         # Calc mean
         for key in top_scores.keys():
             top_scores[key] = np.mean(top_scores[key])

requirements.txt

@@ -2,4 +2,5 @@ gradio
 torch
 transformers
 numpy
-nltk
+nltk
+re