Update app.py

app.py

@@ -72,18 +72,18 @@ def process_entity(batch, model, device):
                         position_list_y=position_list_y)
                         #NOTE: we are ommitting the pseudo_sentence here. Verify that this is correct
 
-    embeddings = outputs.hidden_states[-1].to(device)
+    spaBERT_embedding = outputs.hidden_states[-1].to(device)
 
     # Extract the [CLS] token embedding (first token)
-    embedding = embeddings[:, 0, :].detach()  # [batch_size, hidden_size]
+    spaBERT_embedding = embeddings[:, 0, :].detach()  # [batch_size, hidden_size]
 
     #pivot_token_len = batch['pivot_token_len'].item()
     #pivot_embeddings = embeddings[:, :pivot_token_len, :]
 
     #return pivot_embeddings.cpu().numpy(), input_ids.cpu().numpy()
-    return embedding, input_ids
+    return spaBERT_embedding, input_ids
 
-all_embeddings = []
+spaBERT_embeddings = []
 for i, batch in enumerate(data_loader):
     if i >= 2:  # Stop after processing 3 batches
         break
@@ -105,8 +105,8 @@ def get_bert_embedding(review_text):
         outputs = bert_model(**inputs)
 
     # Extract embeddings from the last hidden state
-    embeddings = outputs.last_hidden_state[:, 0, :].detach()     #CLS Token
-    return embeddings
+    bert_embedding = outputs.last_hidden_state[:, 0, :].detach()     #CLS Token
+    return bert_embedding
 
 
 
@@ -160,7 +160,7 @@ if st.button("Highlight Geo-Entities"):
         st.write("Embeddings:", bert_embedding)
 
         #combine the embeddings (NOTE: come back and update after testing)
-        combined_embedding = torch.cat((bert_embedding,all_embeddings[0]),dim=-1)
+        combined_embedding = torch.cat((bert_embedding,spaBERT_embeddings[0]),dim=-1)
         st.write("Concatenated Embedding Shape:", combined_embedding.shape)
         st.write("Concatenated Embedding:", combined_embedding)