Improve caching

app.py

@@ -60,11 +60,12 @@ def get_word_idx(sent: str, word: str):
     return l.index(word)
 
 
-def get_embedding(tokenizer, model, sent, word, layers=None)-> torch.Tensor:
+@st.cache
+def get_embedding(sent, word, layers=None):
     """Get a word vector by first tokenizing the input sentence, getting all token idxs
         that make up the word of interest, and then `get_hidden_states`."""
     layers = [-4, -3, -2, -1] if layers is None else layers
-
+    tokenizer, model = get_models()
     encoded = tokenizer.encode_plus(sent, return_tensors="pt")
 
     idx = get_word_idx(sent, word)
@@ -81,7 +82,7 @@ def get_embedding(tokenizer, model, sent, word, layers=None)-> torch.Tensor:
     # Only select the tokens that constitute the requested word
     word_tokens_output = output[token_ids_word]
 
-    return word_tokens_output.mean(dim=0)
+    return word_tokens_output.mean(dim=0).numpy()
 
 
 
@@ -100,9 +101,15 @@ for k in verses:
 print(books)
 
 all_defs = {k: f"{k} - {strongs_defs[k]}" for k in strongs_defs}
-option1 = st.multiselect('Select Strongs numbers for first concept', all_defs.keys(),  ['0025', '0026'], format_func=lambda x: strongs_defs[x])
+
+def format_strong(number):
+    return f"{number} - {strongs_defs[number]}"
+
+option1 = st.multiselect('Select Strongs numbers for first concept', all_defs.keys(),  ['0025', '0026'], format_func=format_strong)
 option2 = st.multiselect('Select Strongs numbers for second concept', all_defs.keys(),
-                         ["5368", "5360", "5363", "5362", "5361", "5366", "5365", "5377"], format_func=lambda x: strongs_defs[x])
+                         ["5368", "5360", "5363", "5362", "5361", "5366", "5377"], format_func=format_strong)
+
+
 @st.cache(allow_output_mutation=True)
 def get_models():
     tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
@@ -113,7 +120,6 @@ def get_models():
 @st.cache
 def get_all_embeddings(greek_words):
     embeddings = []
-    tokenizer, model = get_models()
 
     for word in greek_words:
         for number in greek_words[word]:
@@ -122,16 +128,15 @@ def get_all_embeddings(greek_words):
                 for verse, idx in strongs_tags[number]:
                     if verse in verses:
                         text = verses[verse]
-                        print(text, idx)
                         words = [x for x in re.split('([ \'])', text) if x != " " and x != "" and x != "'"]
                         if len(words) <= idx - 1:
                             continue
                         ew = words[idx-1].strip(",.!?;:()\"'-")
-                        print(ew)
-                        emb = get_embedding(tokenizer, model, text, ew).numpy()
-                        embeddings.append((emb, f"{verse} {text}",  ew, book))
+                        emb = get_embedding(text, ew)
+                        embeddings.append((emb, f"{verse} {text}",  gw, book))
     return embeddings
 
+
 def get_book_type(idx):
     if idx < 4:
         return 'Gospels'
@@ -143,11 +148,7 @@ def get_book_type(idx):
         return 'Short lettters'
     return 'Revelation'
 
-strongs_numbers = {
-    "agape": ["0025", "0026"],
-    "phileo": ["5368", "5360", "5363", "5362", "5361", "5366", "5365", "5377"],
-}
-word_list = ["lovers", "loved", "loves", "love", "Love"]
+
 
 embeddings = get_all_embeddings({"concept1": option1, "concept2": option2})