Update app.py

app.py

@@ -1,85 +1,66 @@
 import streamlit as st
-from database_utils import init_db, save_embeddings_to_db, get_all_embeddings, fetch_data_as_csv
+from database_utils import init_db, save_embeddings_to_db, get_all_embeddings, fetch_data_as_csv, clear_all_entries
 from transformers import BertModel, BertTokenizer
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
 import numpy as np
 
 # BERT embeddings function
-def get_bert_embeddings(words):
+def get_bert_embeddings(word):
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
     model = BertModel.from_pretrained('bert-base-uncased')
-    embeddings = []
-    for word in words:
-        inputs = tokenizer(word, return_tensors='pt')
-        outputs = model(**inputs)
-        mean_embedding = outputs.last_hidden_state[0].mean(dim=0).detach().numpy()
-        embeddings.append(mean_embedding)
-    if len(embeddings) > 0:
-        pca = PCA(n_components=3)
-        reduced_embeddings = pca.fit_transform(np.array(embeddings))
-        return reduced_embeddings
-    return []
+    inputs = tokenizer(word, return_tensors='pt')
+    outputs = model(**inputs)
+    mean_embedding = outputs.last_hidden_state[0].mean(dim=0).detach().numpy()
+    pca = PCA(n_components=3)
+    reduced_embedding = pca.fit_transform(np.array([mean_embedding]))  # Embedding needs to be in a 2D array
+    return reduced_embedding[0]  # Return the first and only item
 
-def plot_interactive_bert_embeddings(embeddings, words):
-    if len(words) < 4:
-        st.error("Please provide at least 4 words/phrases for effective visualization.")
-        return None
-    data = []
-    for i, word in enumerate(words):
-        trace = go.Scatter3d(
-            x=[embeddings[i][0]],
-            y=[embeddings[i][1]],
-            z=[embeddings[i][2]],
-            mode='markers+text',
-            text=[word],
-            name=word
-        )
-        data.append(trace)
-    layout = go.Layout(
-        title='3D Scatter Plot of BERT Embeddings',
-        scene=dict(
-            xaxis=dict(title='PCA Component 1'),
-            yaxis=dict(title='PCA Component 2'),
-            zaxis=dict(title='PCA Component 3')
-        ),
-        autosize=False,
-        width=800,
-        height=600
-    )
-    fig = go.Figure(data=data, layout=layout)
-    return fig
-
-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)
+def plot_interactive_bert_embeddings():
+    embeddings, sentences = get_all_embeddings()
+    if len(sentences) >= 3:
+        fig = go.Figure(data=[
+            go.Scatter3d(
+                x=[emb[0]],
+                y=[emb[1]],
+                z=[emb[2]],
+                mode='markers+text',
+                text=[sent],
+                name=sent
+            ) for emb, sent in zip(embeddings, sentences)
+        ], layout=go.Layout(
+            title='3D Scatter Plot of BERT Embeddings',
+            scene=dict(
+                xaxis=dict(title='PCA Component 1'),
+                yaxis=dict(title='PCA Component 2'),
+                zaxis=dict(title='PCA Component 3')
+            ),
+            autosize=False,
+            width=800,
+            height=600
+        ))
+        st.plotly_chart(fig, use_container_width=True)
     else:
-        st.error("Failed to generate embeddings. Ensure your sentences are correctly formatted.")
+        st.error("Please add at least three different words/phrases to visualize.")
 
 def main():
     st.title("BERT Embeddings Visualization - Community Edition")
     
-    init_db_message = init_db()
-    st.text(init_db_message)
+    init_db()
     
-    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)
+    new_word = st.text_input("Enter a new word or phrase:")
+    if st.button("Add Word/Phrase"):
+        if new_word:
+            embedding = get_bert_embeddings(new_word)
+            save_embeddings_to_db(new_word, embedding)
+            st.success(f"Added: {new_word}")
+
+    if st.button("Visualize Embeddings"):
+        plot_interactive_bert_embeddings()
 
-    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 st.button("Clear All Entries"):
+        clear_all_entries()  # This function needs to be defined in database_utils.py to delete all records
+        st.success("All entries have been cleared.")
 
 if __name__ == "__main__":
     main()