address #226

audiolm_pytorch/hubert_kmeans.py

@@ -2,19 +2,32 @@
 
 import torch
 from torch import nn, einsum
-from einops import rearrange, repeat, pack, unpack
-
-import joblib
-
-import fairseq
 
 from torchaudio.functional import resample
 
+from einops import rearrange, repeat, pack, unpack
+
 from audiolm_pytorch.utils import curtail_to_multiple
 
+# suppress a few warnings
+
+def noop(*args, **kwargs):
+    pass
+
+import warnings
 import logging
+
 logging.root.setLevel(logging.ERROR)
 
+warnings.warn = noop
+
+# import fairseq and joblib for hubert
+
+import joblib
+import fairseq
+
+# helper functions
+
 def exists(val):
     return val is not None
 
@@ -37,9 +50,6 @@ def __init__(
     ):
         super().__init__()
 
-        import sklearn
-        assert sklearn.__version__ == '0.24.0', 'scikit-learn needs to be exactly 0.24.0 - please install the correct version by running `pip install scikit-learn==0.24.0`'
-
         self.target_sample_hz = target_sample_hz
         self.seq_len_multiple_of = seq_len_multiple_of
         self.output_layer = output_layer

audiolm_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.4.0'
+__version__ = '1.4.1'