Update app.py

app.py

@@ -3,14 +3,13 @@ import plotly.express as px
 import pandas as pd
 import random
 import logging
-from umap import UMAP
 from sentence_transformers import SentenceTransformer, util
 from datasets import load_dataset
 
 
 @st.cache_resource
-def load_model():
-    return SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
+def load_model(name):
+    return SentenceTransformer(name)
 
 
 @st.cache_data
@@ -24,23 +23,20 @@ def choose_secret_word():
     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=1000))
-    return umap_3d
-
-
 all_words = load_words_dataset()
 
-model = load_model()
 
-umap_3d = prepare_umap()
+model_names = [
+    'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2',
+    'BAAI/bge-small-en-v1.5'
+]
 
+models = {
+    name: load_model(name) for name in model_names
+}
 
-secret_word =choose_secret_word()
-secret_embedding = model.encode(secret_word.lower().strip())
+secret_word =choose_secret_word().lower().strip()
+secret_embedding = [models[name].encode(secret_word) for name in model_names]
 
 print("Secret word ", secret_word)
 
@@ -48,27 +44,6 @@ print("Secret word ", secret_word)
 if 'words' not in st.session_state:
     st.session_state['words'] = []
 
-if 'words_umap_df' not in st.session_state:
-    words_umap_df = pd.DataFrame({
-        "x": [],
-        "y": [],
-        "z": [],
-        "similarity": [],
-        "s": [],
-        "l": [],
-    })
-    st.session_state['words_umap_df'] = 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"
-    }
-    st.session_state['words_umap_df'] = words_umap_df
-
 
 
 
@@ -80,32 +55,15 @@ 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.lower().strip())
-        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 = st.session_state['words_umap_df']
-        words_umap_df.loc[len(words_umap_df)] = {
-            "x": pt[0],
-            "y": pt[1],
-            "z": pt[2],
-            "similarity": similarity,
-            "s": 3,
-            "l": str(word)
-        }
-        st.session_state['words_umap_df'] = words_umap_df
+        word_embedding = [models[name].encode(word.lower().strip()) for name in model_names]
+        similarities = [util.pytorch_cos_sim(secret_embedding[i], word_embedding[i]).cpu().numpy()[0][0] for i, name in enumerate(model_names)]
+        st.session_state['words'].append([str(word)] + similarities))
+
+
 
 words_df = pd.DataFrame(
     st.session_state['words'],
-    columns=["word", "similarity"]
-).sort_values(by=["similarity"], ascending=False)
+    columns=["word"] + ["Similarity for " + name for name in model_names]
+).sort_values(by=["sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"], ascending=False)
 st.dataframe(words_df, use_container_width=True)
 
-
-words_umap_df = st.session_state['words_umap_df']
-fig_3d = px.scatter_3d(words_umap_df, 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)