Update app.py

app.py

@@ -14,9 +14,9 @@ def get_bert_embeddings(words):
     for word in words:
         inputs = tokenizer(word, return_tensors='pt')
         outputs = model(**inputs)
-        # Use the [CLS] token's embedding
-        cls_embedding = outputs.last_hidden_state[0][0].detach().numpy()
-        embeddings.append(cls_embedding)
+        # Calculate mean of embeddings across all tokens in the phrase
+        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)
@@ -25,6 +25,7 @@ def get_bert_embeddings(words):
     return []
 
 
+
 # Plotly plotting function
 def plot_interactive_bert_embeddings(embeddings, words):
     if len(words) < 4: