Merge pull request #175 from wellcometrust/feature/bert_vectorizer_pr…

…ogress_bar new: Add progress bar to transform method
wellcometrust · Nov 18, 2020 · adeecc5 · adeecc5
2 parents 16754b0 + afaaa5b
commit adeecc5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -54,7 +54,8 @@
         'hdbscan',
         'cython',
         'flake8',
-        'black'
+        'black',
+        'tqdm'
     ],
     extras_require={
         'deep-learning': [

diff --git a/wellcomeml/ml/bert_vectorizer.py b/wellcomeml/ml/bert_vectorizer.py
@@ -7,6 +7,7 @@
 from sklearn.base import BaseEstimator, TransformerMixin
 import numpy as np
 import torch
+import tqdm
 
 from wellcomeml.utils import check_cache_and_download
 from wellcomeml.logger import logger
@@ -71,7 +72,7 @@ def bert_embedding(self, x):
         return embedded_x.cpu().numpy().flatten()
 
     def transform(self, X, *_):
-        return np.array([self.bert_embedding(x) for x in X])
+        return np.array([self.bert_embedding(x) for x in tqdm.tqdm(X)])
 
     def fit(self, *_):
         model_name = (