Update app.py

app.py

@@ -84,9 +84,11 @@ def process_entity(batch, model, device):
     return embedding.cpu().numpy(), input_ids.cpu().numpy()
 
 all_embeddings = []
-for batch in (data_loader):
-  embeddings, input_ids = process_entity(batch, spaBERT_model, device)
-  all_embeddings.append(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)
 
 st.write("SpaBERT Embedding shape:", all_embeddings[0].shape)
 st.write("SpaBERT Embedding:", all_embeddings[0])