Update app.py

app.py

@@ -4,7 +4,27 @@ import torch
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
 
-# BERT Embeddings
+# BERT embeddings function
+def get_bert_embeddings(words):
+    # Load pre-trained BERT model and tokenizer
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    model = BertModel.from_pretrained('bert-base-uncased')
+
+    embeddings = []
+    
+    # Extract embeddings
+    for word in words:
+        inputs = tokenizer(word, return_tensors='pt')
+        outputs = model(**inputs)
+        embeddings.append(outputs.last_hidden_state[0][0].detach().numpy())
+
+    # Reduce dimensions to 3 using PCA
+    pca = PCA(n_components=3)
+    reduced_embeddings = pca.fit_transform(embeddings)
+
+    return reduced_embeddings
+
+# 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.")
@@ -49,7 +69,8 @@ def main():
         with st.spinner('Generating embeddings...'):
             embeddings = get_bert_embeddings(words)
             fig = plot_interactive_bert_embeddings(embeddings, words)
-            st.plotly_chart(fig, use_container_width=True)
+            if fig is not None:  # Only plot if the figure is not None
+                st.plotly_chart(fig, use_container_width=True)
 
 if __name__ == "__main__":
-    main()
+    main()