Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from database_utils import init_db, save_embeddings_to_db, get_all_embeddings
+from database_utils import init_db, save_embeddings_to_db, get_all_embeddings, fetch_data_as_csv
 from transformers import BertModel, BertTokenizer
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
@@ -13,7 +13,6 @@ def get_bert_embeddings(words):
     for word in words:
         inputs = tokenizer(word, return_tensors='pt')
         outputs = model(**inputs)
-        # Average the embeddings of all tokens for the word/phrase
         mean_embedding = outputs.last_hidden_state[0].mean(dim=0).detach().numpy()
         embeddings.append(mean_embedding)
     if len(embeddings) > 0:
@@ -22,7 +21,6 @@ def get_bert_embeddings(words):
         return reduced_embeddings
     return []
 
-# Plotly plotting function
 def plot_interactive_bert_embeddings(embeddings, words):
     if len(words) < 4:
         st.error("Please provide at least 4 words/phrases for effective visualization.")
@@ -52,29 +50,36 @@ def plot_interactive_bert_embeddings(embeddings, words):
     fig = go.Figure(data=data, layout=layout)
     return fig
 
-def main():
-    st.title("BERT Embeddings Visualization - Community Edition")
-
-    # Button to initialize the database
-    if st.button("Initialize Database"):
-        msg = init_db()
-        st.success(msg)
-
-    # Text input for new sentence
-    new_sentence = st.text_input("Enter a new sentence:")
-    if st.button("Add and Visualize Sentence"):
-        if new_sentence:
-            embeddings = get_bert_embeddings([new_sentence])
-            if embeddings.size > 0:
-                save_embeddings_to_db(new_sentence, embeddings[0])
-                st.success("Sentence added and embedding saved!")
-
-    # Button to display all embeddings
-    if st.button("Show All Embeddings"):
+def process_and_visualize_words(words):
+    embeddings = get_bert_embeddings(words)
+    if embeddings.size > 0:
+        for word, embedding in zip(words, embeddings):
+            save_embeddings_to_db(word, embedding)
+        st.success("Sentences added and embeddings saved!")
         embeddings, sentences = get_all_embeddings()
         fig = plot_interactive_bert_embeddings(np.vstack(embeddings), sentences)
         if fig is not None:
             st.plotly_chart(fig, use_container_width=True)
+    else:
+        st.error("Failed to generate embeddings. Ensure your sentences are correctly formatted.")
+
+def main():
+    st.title("BERT Embeddings Visualization - Community Edition")
+    
+    init_db_message = init_db()
+    st.text(init_db_message)
+    
+    new_sentences = st.text_input("Enter at least three words/phrases, comma-separated:", "Example: apple, banana, orange")
+    if st.button("Add and Visualize Sentences"):
+        words = [word.strip() for word in new_sentences.split(',')]
+        if len(words) < 3:
+            st.error("Please enter at least three words/phrases separated by commas.")
+        else:
+            process_and_visualize_words(words)
+
+    if st.button("Download Database as CSV"):
+        csv = fetch_data_as_csv()
+        st.download_button(label="Download CSV", data=csv, file_name='embeddings.csv', mime='text/csv')
 
 if __name__ == "__main__":
     main()