Update app.py

app.py

@@ -18,23 +18,28 @@ def load_words_dataset():
     dataset = load_dataset("marksverdhei/wordnet-definitions-en-2021", split="train")
     return dataset["Word"]
 
+@st.cache_data
+def choose_secret_word():
+    all_words = load_words_dataset()
+    return random.choice(all_words)
+
 
-# @st.cache_resource
-# def prepare_umap():
-#     all_enc = model.encode(all_words)
-#     umap_3d = UMAP(n_components=3, init='random', random_state=0)
-#     proj_3d = umap_3d.fit_transform(random.sample(all_enc.tolist(), k=2000))
-#     return umap_3d
+@st.cache_resource
+def prepare_umap():
+    all_enc = model.encode(all_words)
+    umap_3d = UMAP(n_components=3, init='random', random_state=0)
+    proj_3d = umap_3d.fit_transform(random.sample(all_enc.tolist(), k=1000))
+    return umap_3d
 
 
 all_words = load_words_dataset()
 
 model = load_model()
 
-#umap_3d = prepare_umap()
+umap_3d = prepare_umap()
 
 
-secret_word = random.choice(all_words)
+secret_word =choose_secret_word()
 secret_embedding = model.encode(secret_word)
 
 print("Secret word ", secret_word)
@@ -52,8 +57,7 @@ if 'words_umap_df' not in st.session_state:
         "s": [],
         "l": [],
     })
-    #secret_embedding_3d = umap_3d.transform([secret_embedding])[0]
-    secret_embedding_3d = [0, 1, 2]
+    secret_embedding_3d = umap_3d.transform([secret_embedding])[0]
     words_umap_df.loc[len(words_umap_df)] = {
         "x": secret_embedding_3d[0],
         "y": secret_embedding_3d[1],
@@ -82,8 +86,7 @@ if st.button("Guess") or word:
         ).cpu().numpy()[0][0]
         st.session_state['words'].append((str(word), similarity))
 
-        #pt = umap_3d.transform([word_embedding])[0]
-        pt = [0, 1, 2]
+        pt = umap_3d.transform([word_embedding])[0]
         words_umap_df = st.session_state['words_umap_df']
         words_umap_df.loc[len(words_umap_df)] = {
             "x": pt[0],