Update app.py

app.py

@@ -11,19 +11,20 @@ def get_bert_embeddings(words):
     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())
+        # Use the [CLS] token's embedding
+        cls_embedding = outputs.last_hidden_state[0][0].detach().numpy()
+        embeddings.append(cls_embedding)
 
-    # Reduce dimensions to 3 using PCA
     if len(embeddings) > 0:
         pca = PCA(n_components=3)
         reduced_embeddings = pca.fit_transform(np.array(embeddings))
         return reduced_embeddings
     return []
 
+
 # Plotly plotting function
 def plot_interactive_bert_embeddings(embeddings, words):
     if len(words) < 4: