Update app.py

app.py

@@ -1,5 +1,8 @@
 import streamlit as st
+import plotly.express as px
 import pandas as pd
+import random
+from umap import UMAP
 from sentence_transformers import SentenceTransformer, util
 from datasets import load_dataset
 
@@ -8,14 +11,58 @@ from datasets import load_dataset
 def load_model():
     return SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
 
+
+@st.cache_data
+def load_words_dataset():
+    dataset = load_dataset("marksverdhei/wordnet-definitions-en-2021", split="train")
+    return dataset["Word"]
+
+
+@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(all_enc)
+    return umap_3d
+
+
+all_words = load_words_dataset()
+
 model = load_model()
 
-secret_word = "нос"
-secred_embedding = model.encode(secret_word)
+umap_3d = prepare_umap()
+
+
+secret_word = random.choice(all_words)
+secret_embedding = model.encode(secret_word)
+
 
 if 'words' not in st.session_state:
     st.session_state['words'] = []
 
+if 'words_umap_df' not in st.session_state:
+    st.session_state['words_umap_df'] = pd.DataFrame({
+        "x": [],
+        "y": [],
+        "z": [],
+        "similarity": [],
+        "s": [],
+        "l": [],
+    })
+    words_umap_df = st.session_state['words_umap_df']
+
+    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],
+        "z": secret_embedding_3d[2],
+        "similarity": 1,
+        "s": 10,
+        "l": "Secret word"
+    }
+
+words_umap_df = st.session_state['words_umap_df']
+
 st.write('Try to guess a secret word by semantic similarity')
 
 word = st.text_input("Input a word")
@@ -25,8 +72,21 @@ used_words = [w for w, s in st.session_state['words']]
 if st.button("Guess") or word:
     if word not in used_words:
         word_embedding = model.encode(word)
-        similarity = util.pytorch_cos_sim(secred_embedding, word_embedding).cpu().numpy()[0][0]
-        st.session_state['words'].append((word, similarity))
+        similarity = util.pytorch_cos_sim(
+            secret_embedding,
+            word_embedding
+        ).cpu().numpy()[0][0]
+        st.session_state['words'].append((str(word), similarity))
+
+        pt = umap_3d.transform([word_embedding])[0]
+        words_umap_df.loc[len(words_umap_df)] = {
+            "x": pt[0],
+            "y": pt[1],
+            "z": pt[2],
+            "similarity": similarity,
+            "s": 3,
+            "l": str(word)
+        }
 
 words_df = pd.DataFrame(
     st.session_state['words'],
@@ -35,11 +95,5 @@ words_df = pd.DataFrame(
 st.dataframe(words_df)
 
 
-
-@st.cache_data
-def load_words_dataset():
-    dataset = load_dataset("marksverdhei/wordnet-definitions-en-2021", split="train")
-    return dataset["Word"]
-
-all_words = load_words_dataset()
-st.write(all_words)
+fig_3d = px.scatter_3d(word_points, x="x", y="y", z="z", color="similarity", hover_name="l", hover_data={"x": False, "y": False, "z": False, "s": False}, size="s", size_max=10, range_color=(0,1))
+st.plotly_chart(fig_3d, theme="streamlit", use_container_width=True)