added tqdm for progess bar

stair-lab · Jul 15, 2024 · 7a40a08 · 7a40a08
1 parent 2d4b629
commit 7a40a08
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/embed_text_package/embed_text.py b/src/embed_text_package/embed_text.py
@@ -3,6 +3,7 @@
 """
 
 import torch
+from tqdm import tqdm
 
 
 def get_embeddings(sentence_batches: list, model, tokenizer):
@@ -31,7 +32,8 @@ def get_embeddings(sentence_batches: list, model, tokenizer):
 
     emb_batches = []
 
-    for batch in sentence_batches:
+    for batch in tqdm(sentence_batches,
+                      ascii=True, desc="Embedding Batches..."):
         batch_emb = []
         for sentence in batch:
             # 1) Get Tokens of sentence