Update app.py

app.py

@@ -7,8 +7,21 @@ from transformers import AutoTokenizer
 
 tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased")
 input_embeddings = np.load("bert_input_embeddings.npy")
-index = IndexFlatL2(input_embeddings.shape[-1])  
-index.add(input_embeddings)
+unnormalized_input_embeddings = np.load("unnormalized_bert_input_embeddings.npy")
+
+index_L2 = IndexFlatL2(input_embeddings.shape[-1])  
+index_L2.add(input_embeddings)
+
+index_IP = IndexFlatIP(input_embeddings.shape[-1])  
+index_IP.add(input_embeddings)
+
+index_L2_unnormalized = IndexFlatL2(unnormalized_input_embeddings.shape[-1])  
+index_L2_unnormalized.add(unnormalized_input_embeddings)
+
+index_IP_unnormalized = IndexFlatIP(unnormalized_input_embeddings.shape[-1])  
+index_IP_unnormalized.add(unnormalized_input_embeddings)
+
+
 vocab = {v:k for k,v in tokenizer.vocab.items()}
 lookup_table = pd.Series(vocab).sort_index()
 
@@ -18,9 +31,9 @@ def get_first_subword(word):
     except:
         return tokenizer(word, add_special_tokens=False)['input_ids'][0]
 
-def search(token_to_lookup, num_neighbors=50):
+def search(token_to_lookup, num_neighbors=200):
     i = get_first_subword(token_to_lookup)
-    _ , I = index.search(input_embeddings[i:i+1], num_neighbors)
+    _ , I = index_L2_unnormalized.search(unnormalized_input_embeddings[i:i+1], num_neighbors)
     hits = lookup_table.take(I[0])
     results = hits.values[1:]
     return [r for r in results if not "##" in r], [[r for r in results if "##" in r]]