Update app.py

app.py

@@ -15,8 +15,8 @@ def get_bert_embeddings(words):
         inputs = tokenizer(word, return_tensors='pt')
         outputs = model(**inputs)
         mean_embedding = outputs.last_hidden_state.mean(dim=1).detach().numpy()
-        embeddings.append(mean_embedding[0])
-    return embeddings
+        embeddings.append(mean_embedding[0])  # Append the 1D embedding
+    return np.array(embeddings)
 
 def plot_interactive_bert_embeddings(embeddings, words):
     if len(words) >= 3:  # Ensure there are at least 3 words for 3D PCA
@@ -28,7 +28,7 @@ def plot_interactive_bert_embeddings(embeddings, words):
                 y=[emb[1]],
                 z=[emb[2]],
                 mode='markers+text',
-                text=word,
+                text=[word],
                 name=word
             ) for emb, word in zip(reduced_embeddings, words)
         ], layout=go.Layout(
@@ -51,7 +51,7 @@ def main():
 
     default_words = ["apple", "rocket", "philosophy"]
     if "words" not in st.session_state:
-        st.session_state.words = default_words
+        st.session_state.words = default_words[:]
         init_db()  # Initialize the database
         for word in default_words:
             embedding = get_bert_embeddings([word])[0]
@@ -83,7 +83,9 @@ def main():
         st.download_button(label="Download CSV", data=csv, file_name='embeddings.csv', mime='text/csv')
 
     embeddings, words = get_all_embeddings()
-    plot_interactive_bert_embeddings(embeddings, words)
+    embeddings = np.array(embeddings)
+    if embeddings.size > 0:
+        plot_interactive_bert_embeddings(embeddings, words)
 
 if __name__ == "__main__":
     main()