Update app.py

app.py

@@ -6,20 +6,14 @@ from transformers import AutoTokenizer
 
 
 tokenizer = AutoTokenizer.from_pretrained("bert-large-uncased")
-input_embeddings = np.load("bert_input_embeddings.npy")
+normalized_input_embeddings = np.load("normalized_bert_input_embeddings.npy")
 unnormalized_input_embeddings = np.load("unnormalized_bert_input_embeddings.npy")
 
-index_L2 = IndexFlatL2(input_embeddings.shape[-1])  
-index_L2.add(input_embeddings)
+index_L2 = IndexFlatL2(unnormalized_input_embeddings.shape[-1])  
+index_L2.add(unnormalized_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)
+index_IP_normalized = IndexFlatIP(normalized_input_embeddings.shape[-1])  
+index_IP_normalized.add(normalized_input_embeddings)
 
 
 vocab = {v:k for k,v in tokenizer.vocab.items()}
@@ -33,7 +27,7 @@ def get_first_subword(word):
 
 def search(token_to_lookup, num_neighbors):
     i = get_first_subword(token_to_lookup)
-    _ , I = index_L2_unnormalized.search(unnormalized_input_embeddings[i:i+1], num_neighbors)
+    _ , I = index_IP_normalized.search(normalized_input_embeddings[i:i+1], num_neighbors)
     hits = lookup_table.take(I[0])
     results = hits.values[1:]
     results = [r for r in results if not "[unused" in r]