From 9193acdd29b2985596e050aa0ce43ff313cafa2e Mon Sep 17 00:00:00 2001
From: Michael-Danziger <michael.danziger@ibm.com>
Date: Wed, 18 Sep 2024 12:24:46 +0300
Subject: [PATCH] Add tqdm for encoders

It can take a long time with deep learning models, tqdm is really useful
---
 gene_benchmark/encoder.py | 3 ++-
 pyproject.toml            | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/gene_benchmark/encoder.py b/gene_benchmark/encoder.py
index 05375f4..87c5645 100644
--- a/gene_benchmark/encoder.py
+++ b/gene_benchmark/encoder.py
@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import torch
+import tqdm
 from sentence_transformers import SentenceTransformer
 from transformers import AutoModel, AutoTokenizer
 from transformers.models.bert.configuration_bert import BertConfig
@@ -442,7 +443,7 @@ def __init__(
 
     def _get_encoding(self, entities, **kwargs):
         vec_list = []
-        for ent in entities:
+        for ent in tqdm.tqdm(entities):
             inputs = self.tokenizer(ent, return_tensors="pt")["input_ids"]
             hidden_states = self.encoder(inputs)[0]
             vec_list.append(torch.mean(hidden_states[0], dim=0).detach())
diff --git a/pyproject.toml b/pyproject.toml
index c4539d7..63036bd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,7 +25,8 @@ dependencies = [
 "sentence_transformers",
 "scikit-learn",
 "click",
-"einops"
+"einops",
+"tqdm"
 ]
 
 [project.optional-dependencies]