Update app.py

app.py

@@ -1,81 +1,109 @@
 import streamlit as st
-from transformers import BertModel, BertTokenizer
+from transformers import BertModel, BertTokenizer, RobertaModel, RobertaTokenizer
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
 import numpy as np
 from database_utils import init_db, save_embeddings_to_db, get_all_embeddings, clear_all_entries, fetch_data_as_csv
 
-# Initialize BERT model and tokenizer
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-model = BertModel.from_pretrained('bert-base-uncased')
+@st.cache_resource
+def load_model(model_name):
+    if model_name == "BERT":
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        model = BertModel.from_pretrained('bert-base-uncased')
+    elif model_name == "RoBERTa":
+        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+        model = RobertaModel.from_pretrained('roberta-base')
+    return tokenizer, model
 
-def get_bert_embeddings(words):
+def get_embeddings(words, tokenizer, model):
     embeddings = []
     for word in words:
-        inputs = tokenizer(word, return_tensors='pt')
+        inputs = tokenizer(word, return_tensors='pt', padding=True, truncation=True)
         outputs = model(**inputs)
         mean_embedding = outputs.last_hidden_state.mean(dim=1).detach().numpy()
-        embeddings.append(mean_embedding[0])  # Append the 1D embedding
+        embeddings.append(mean_embedding[0])
     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
-        pca = PCA(n_components=3)
+def plot_interactive_embeddings(embeddings, words):
+    if len(words) >= 2:
+        pca = PCA(n_components=min(3, len(words)))
         reduced_embeddings = pca.fit_transform(embeddings)
-        fig = go.Figure(data=[
-            go.Scatter3d(
-                x=[emb[0]],
-                y=[emb[1]],
-                z=[emb[2]],
-                mode='markers+text',
-                text=[word],
-                name=word
-            ) for emb, word in zip(reduced_embeddings, words)
-        ], 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
-        ))
+        
+        if len(words) == 2:
+            fig = go.Figure(data=[
+                go.Scatter(
+                    x=[emb[0]],
+                    y=[emb[1]],
+                    mode='markers+text',
+                    text=[word],
+                    name=word
+                ) for emb, word in zip(reduced_embeddings, words)
+            ])
+            fig.update_layout(
+                title='2D Scatter Plot of Embeddings',
+                xaxis_title='PCA Component 1',
+                yaxis_title='PCA Component 2'
+            )
+        else:
+            fig = go.Figure(data=[
+                go.Scatter3d(
+                    x=[emb[0]],
+                    y=[emb[1]],
+                    z=[emb[2]],
+                    mode='markers+text',
+                    text=[word],
+                    name=word
+                ) for emb, word in zip(reduced_embeddings, words)
+            ])
+            fig.update_layout(
+                title='3D Scatter Plot of Embeddings',
+                scene=dict(
+                    xaxis_title='PCA Component 1',
+                    yaxis_title='PCA Component 2',
+                    zaxis_title='PCA Component 3'
+                )
+            )
+        
+        fig.update_layout(autosize=False, width=800, height=600)
         st.plotly_chart(fig, use_container_width=True)
     else:
-        st.error("Please add more words to visualize. A minimum of three is required.")
+        st.error("Please add at least one more word to visualize.")
 
 def main():
-    st.title("BERT Embeddings Visualization")
+    st.title("Language Model Embeddings Visualization")
 
-    default_words = ["apple", "rocket", "philosophy"]
-    if "words" not in st.session_state:
-        st.session_state.words = default_words[:]
-        init_db()  # Initialize the database
-        for word in default_words:
-            embedding = get_bert_embeddings([word])[0]
-            save_embeddings_to_db(word, embedding)
+    model_choice = st.selectbox("Choose a model:", ["BERT", "RoBERTa"])
+    tokenizer, model = load_model(model_choice)
 
-    if st.button("Reset to Default Words"):
+    default_word = "example"
+    if "words" not in st.session_state or "model" not in st.session_state:
+        st.session_state.words = [default_word]
+        st.session_state.model = model_choice
+        init_db()
+        embedding = get_embeddings([default_word], tokenizer, model)[0]
+        save_embeddings_to_db(default_word, embedding)
+    elif st.session_state.model != model_choice:
+        st.session_state.words = [default_word]
+        st.session_state.model = model_choice
         clear_all_entries()
-        st.session_state.words = default_words[:]
-        for word in default_words:
-            embedding = get_bert_embeddings([word])[0]
-            save_embeddings_to_db(word, embedding)
-        st.experimental_rerun()
+        embedding = get_embeddings([default_word], tokenizer, model)[0]
+        save_embeddings_to_db(default_word, embedding)
+
+    st.write(f"Current words ({model_choice}):", ", ".join(st.session_state.words))
 
-    new_word = st.text_input("Enter a new word or phrase:")
+    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])[0]
+            embedding = get_embeddings([new_word], tokenizer, model)[0]
             save_embeddings_to_db(new_word, embedding)
             st.session_state.words.append(new_word)
             st.experimental_rerun()
 
     if st.button("Clear All Entries"):
         clear_all_entries()
-        st.session_state.words = default_words[:]
+        st.session_state.words = [default_word]
+        embedding = get_embeddings([default_word], tokenizer, model)[0]
+        save_embeddings_to_db(default_word, embedding)
         st.experimental_rerun()
 
     if st.button("Download Database as CSV"):
@@ -85,7 +113,7 @@ def main():
     embeddings, words = get_all_embeddings()
     if len(embeddings) > 0:
         embeddings = np.array(embeddings)
-        plot_interactive_bert_embeddings(embeddings, words)
+        plot_interactive_embeddings(embeddings, words)
 
 if __name__ == "__main__":
-    main()
+    main()