Update app.py

app.py

@@ -75,7 +75,7 @@ def process_entity(batch, model, device):
     spaBERT_embedding = outputs.hidden_states[-1].to(device)
 
     # Extract the [CLS] token embedding (first token)
-    spaBERT_embedding = embeddings[:, 0, :].detach()  # [batch_size, hidden_size]
+    spaBERT_embedding = spaBERT_embedding[:, 0, :].detach()  # [batch_size, hidden_size]
 
     #pivot_token_len = batch['pivot_token_len'].item()
     #pivot_embeddings = embeddings[:, :pivot_token_len, :]
@@ -87,11 +87,11 @@ spaBERT_embeddings = []
 for i, batch in enumerate(data_loader):
     if i >= 2:  # Stop after processing 3 batches
         break
-    embeddings, input_ids = process_entity(batch, spaBERT_model, device)
-    all_embeddings.append(embeddings)
+    spaBERT_embedding, input_ids = process_entity(batch, spaBERT_model, device)
+    spaBERT_embeddings.append(spaBERT_embedding)
 
-st.write("SpaBERT Embedding shape:", all_embeddings[0].shape)
-st.write("SpaBERT Embedding:", all_embeddings[0])
+st.write("SpaBERT Embedding shape:", spaBERT_embedding[0].shape)
+st.write("SpaBERT Embedding:", spaBERT_embedding[0])