Merge pull request #175 from wellcometrust/feature/bert_vectorizer_progress_bar

new: Add progress bar to transform method

Author: nsorros

setup.py

@@ -54,7 +54,8 @@
         'hdbscan',
         'cython',
         'flake8',
-        'black'
+        'black',
+        'tqdm'
     ],
     extras_require={
         'deep-learning': [

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 = (