Update app.py

app.py

@@ -4,30 +4,13 @@ import torch
 from sklearn.decomposition import PCA
 import plotly.graph_objs as go
 
-# 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
+# BERT Embeddings
 def plot_interactive_bert_embeddings(embeddings, words):
-    data = []
+    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]], 
@@ -45,7 +28,10 @@ def plot_interactive_bert_embeddings(embeddings, words):
             xaxis=dict(title='PCA Component 1'),
             yaxis=dict(title='PCA Component 2'),
             zaxis=dict(title='PCA Component 3')
-        )
+        ),
+        autosize=False,
+        width=800,  # Width of the plot
+        height=600  # Height of the plot
     )
 
     fig = go.Figure(data=data, layout=layout)