nkululeko 0.81.4__py3-none-any.whl → 0.81.7__py3-none-any.whl

nkululeko/feat_extract/feats_snr.py

@@ -1,14 +1,17 @@
-""" feats_snr.py
-    Estimate snr (signal to noise ratio as acoustic features)
+""" feats_snr.py is to estimate snr.
+
+SNR (signal to noise ratio) is extracted as acoustic features.
 """
 import os
-from tqdm import tqdm
-import pandas as pd
+
 import audiofile
+import pandas as pd
+from tqdm import tqdm
+
 import nkululeko.glob_conf as glob_conf
-from nkululeko.utils.util import Util
-from nkululeko.feat_extract.featureset import Featureset
 from nkululeko.autopredict.estimate_snr import SNREstimator
+from nkululeko.feat_extract.featureset import Featureset
+from nkululeko.utils.util import Util
 
 
 class SNRSet(Featureset):
@@ -16,14 +19,17 @@ class SNRSet(Featureset):
 
     def __init__(self, name, data_df):
         """Constructor."""
+
         super().__init__(name, data_df)
 
     def extract(self):
         """Estimate the features or load them from disk if present."""
+
         store = self.util.get_path("store")
         store_format = self.util.config_val("FEATS", "store_format", "pkl")
         storage = f"{store}{self.name}.{store_format}"
-        extract = self.util.config_val("FEATS", "needs_feature_extraction", False)
+        extract = self.util.config_val(
+            "FEATS", "needs_feature_extraction", False)
         no_reuse = eval(self.util.config_val("FEATS", "no_reuse", "False"))
         if extract or no_reuse or not os.path.isfile(storage):
             self.util.debug("estimating SNR, this might take a while...")
@@ -40,7 +46,8 @@ class SNRSet(Featureset):
                 snr = self.get_snr(signal[0], sampling_rate)
                 snr_series[idx] = snr
             print("")
-            self.df = pd.DataFrame(snr_series.values.tolist(), index=self.data_df.index)
+            self.df = pd.DataFrame(
+                snr_series.values.tolist(), index=self.data_df.index)
             self.df.columns = ["snr"]
             self.util.write_store(self.df, storage, store_format)
             try:
@@ -53,10 +60,11 @@ class SNRSet(Featureset):
 
     def get_snr(self, signal, sampling_rate):
         r"""Estimate SNR from raw audio signal.
+
         Args:
             signal: audio signal
             sampling_rate: sample rate
-        Returns
+        Returns:
             snr: estimated signal to noise ratio
         """
         snr_estimator = SNREstimator(signal, sampling_rate)

nkululeko/feat_extract/feats_spectra.py

@@ -4,6 +4,7 @@ feats_spectra.py
 Inspired by code from Su Lei
 
 """
+
 import os
 import torchaudio
 import torchaudio.transforms as T
@@ -23,9 +24,9 @@ import nkululeko.glob_conf as glob_conf
 
 
 class Spectraloader(Featureset):
-    def __init__(self, name, data_df):
+    def __init__(self, name, data_df, feat_type):
         """Constructor setting the name"""
-        Featureset.__init__(self, name, data_df)
+        super().__init__(name, data_df, feat_type)
         self.sampling_rate = SAMPLING_RATE
         self.num_bands = int(self.util.config_val("FEATS", "fft_nbands", "64"))
         self.win_dur = int(self.util.config_val("FEATS", "fft_win_dur", "25"))

nkululeko/feat_extract/feats_squim.py

@@ -1,41 +1,38 @@
-""" feats_squim.py
-predict SQUIM ( SPEECH QUALITY AND INTELLIGIBILITY
-MEASURES) features
+"""Predict SQUIM ( SPEECH QUALITY AND INTELLIGIBILITY MEASURES) features.
 
-
-    Wideband Perceptual Estimation of Speech Quality (PESQ) [2]
-    Short-Time Objective Intelligibility (STOI) [3]
-    Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) [4]
-
-
-adapted from
+Wideband Perceptual Estimation of Speech Quality (PESQ) [2].
+    Short-Time Objective Intelligibility (STOI) [3].
+    Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) [4].
+Adapted from
 from https://pytorch.org/audio/main/tutorials/squim_tutorial.html#sphx-glr-tutorials-squim-tutorial-py
-paper: https://arxiv.org/pdf/2304.01448.pdf
-
-needs 
+paper: https://arxiv.org/pdf/2304.01448.pdf.
+Needs
 pip uninstall -y torch torchvision torchaudio
 pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu
 
 """
 
 import os
-from tqdm import tqdm
+
 import pandas as pd
 import torch
 import torchaudio
 from torchaudio.pipelines import SQUIM_OBJECTIVE
+from tqdm import tqdm
+
 import audiofile
+
+from nkululeko.feat_extract.featureset import Featureset
 import nkululeko.glob_conf as glob_conf
 from nkululeko.utils.util import Util
-from nkululeko.feat_extract.featureset import Featureset
 
 
-class SQUIMSet(Featureset):
+class SquimSet(Featureset):
     """Class to predict SQUIM features"""
 
-    def __init__(self, name, data_df):
+    def __init__(self, name, data_df, feats_type):
         """Constructor. is_train is needed to distinguish from test/dev sets, because they use the codebook from the training"""
-        super().__init__(name, data_df)
+        super().__init__(name, data_df, feats_type)
         self.device = self.util.config_val("MODEL", "device", "cpu")
         self.model_initialized = False
 

nkululeko/feat_extract/feats_trill.py

@@ -1,4 +1,5 @@
 # feats_trill.py
+import tensorflow_hub as hub
 import os
 import tensorflow as tf
 from numpy.core.numeric import tensordot
@@ -11,7 +12,6 @@ from nkululeko.feat_extract.featureset import Featureset
 
 # Import TF 2.X and make sure we're running eager.
 assert tf.executing_eagerly()
-import tensorflow_hub as hub
 
 
 class TRILLset(Featureset):
@@ -20,7 +20,7 @@ class TRILLset(Featureset):
     """https://ai.googleblog.com/2020/06/improving-speech-representations-and.html"""
 
     # Initialization of the class
-    def __init__(self, name, data_df):
+    def __init__(self, name, data_df, feats_type):
         """
         Initialize the class with name, data and Util instance
         Also loads the model from hub
@@ -31,7 +31,7 @@ class TRILLset(Featureset):
         :type data_df: DataFrame
         :return: None
         """
-        super().__init__(name, data_df)
+        super().__init__(name, data_df, feats_type)
         # Load the model from the configured path
         model_path = self.util.config_val(
             "FEATS",
@@ -39,20 +39,24 @@ class TRILLset(Featureset):
             "https://tfhub.dev/google/nonsemantic-speech-benchmark/trill/3",
         )
         self.module = hub.load(model_path)
+        self.feats_type = feats_type
 
     def extract(self):
         store = self.util.get_path("store")
         storage = f"{store}{self.name}.pkl"
-        extract = self.util.config_val("FEATS", "needs_feature_extraction", False)
+        extract = self.util.config_val(
+            "FEATS", "needs_feature_extraction", False)
         no_reuse = eval(self.util.config_val("FEATS", "no_reuse", "False"))
         if extract or no_reuse or not os.path.isfile(storage):
-            self.util.debug("extracting TRILL embeddings, this might take a while...")
+            self.util.debug(
+                "extracting TRILL embeddings, this might take a while...")
             emb_series = pd.Series(index=self.data_df.index, dtype=object)
             length = len(self.data_df.index)
             for idx, file in enumerate(tqdm(self.data_df.index.get_level_values(0))):
                 emb = self.getEmbeddings(file)
                 emb_series[idx] = emb
-            self.df = pd.DataFrame(emb_series.values.tolist(), index=self.data_df.index)
+            self.df = pd.DataFrame(
+                emb_series.values.tolist(), index=self.data_df.index)
             self.df.to_pickle(storage)
             try:
                 glob_conf.config["DATA"]["needs_feature_extraction"] = "false"

nkululeko/feat_extract/feats_wav2vec2.py

@@ -1,5 +1,11 @@
-# feats_wav2vec2.py
-# feat_types example = wav2vec2-large-robust-ft-swbd-300h
+""" feats_wav2vec2.py
+feat_types example = [wav2vec2-large-robust-ft-swbd-300h,
+wav2vec2-xls-r-2b, wav2vec2-large, wav2vec2-large-xlsr-53, wav2vec2-base]
+
+Complete list: https://huggingface.co/facebook?search_models=wav2vec2
+Currently only supports wav2vec2
+"""
+
 import os
 from tqdm import tqdm
 import pandas as pd
@@ -16,11 +22,11 @@ class Wav2vec2(Featureset):
 
     def __init__(self, name, data_df, feat_type):
         """Constructor. is_train is needed to distinguish from test/dev sets, because they use the codebook from the training"""
-        super().__init__(name, data_df)
+        super().__init__(name, data_df, feat_type)
         cuda = "cuda" if torch.cuda.is_available() else "cpu"
         self.device = self.util.config_val("MODEL", "device", cuda)
         self.model_initialized = False
-        if feat_type == "wav2vec" or feat_type == "wav2vec2":
+        if feat_type == "wav2vec2":
             self.feat_type = "wav2vec2-large-robust-ft-swbd-300h"
         else:
             self.feat_type = feat_type
@@ -33,7 +39,8 @@ class Wav2vec2(Featureset):
         )
         config = transformers.AutoConfig.from_pretrained(model_path)
         layer_num = config.num_hidden_layers
-        hidden_layer = int(self.util.config_val("FEATS", "wav2vec2.layer", "0"))
+        hidden_layer = int(self.util.config_val(
+            "FEATS", "wav2vec2.layer", "0"))
         config.num_hidden_layers = layer_num - hidden_layer
         self.util.debug(f"using hidden layer #{config.num_hidden_layers}")
         self.processor = Wav2Vec2FeatureExtractor.from_pretrained(model_path)
@@ -48,7 +55,8 @@ class Wav2vec2(Featureset):
         """Extract the features or load them from disk if present."""
         store = self.util.get_path("store")
         storage = f"{store}{self.name}.pkl"
-        extract = self.util.config_val("FEATS", "needs_feature_extraction", False)
+        extract = self.util.config_val(
+            "FEATS", "needs_feature_extraction", False)
         no_reuse = eval(self.util.config_val("FEATS", "no_reuse", "False"))
         if extract or no_reuse or not os.path.isfile(storage):
             if not self.model_initialized:
@@ -69,7 +77,8 @@ class Wav2vec2(Featureset):
                 emb = self.get_embeddings(signal, sampling_rate, file)
                 emb_series[idx] = emb
             # print(f"emb_series shape: {emb_series.shape}")
-            self.df = pd.DataFrame(emb_series.values.tolist(), index=self.data_df.index)
+            self.df = pd.DataFrame(
+                emb_series.values.tolist(), index=self.data_df.index)
             # print(f"df shape: {self.df.shape}")
             self.df.to_pickle(storage)
             try:

nkululeko/feat_extract/feats_wavlm.py

@@ -59,10 +59,7 @@ class Wavlm(Featureset):
                     frame_offset=int(start.total_seconds() * 16000),
                     num_frames=int((end - start).total_seconds() * 16000),
                 )
-                if sampling_rate != 16000:
-                    self.util.error(
-                        f"sampling rate should be 16000 but is {sampling_rate}"
-                    )
+                assert sampling_rate == 16000, f"sampling rate should be 16000 but is {sampling_rate}"
                 emb = self.get_embeddings(signal, sampling_rate, file)
                 emb_series.iloc[idx] = emb
             self.df = pd.DataFrame(emb_series.values.tolist(), index=self.data_df.index)

nkululeko/feat_extract/feats_whisper.py

@@ -0,0 +1,110 @@
+# feats_whisper.py
+import os
+
+import pandas as pd
+import torch
+from transformers import AutoFeatureExtractor
+from transformers import WhisperModel
+
+import audeer
+import audiofile
+
+from nkululeko.feat_extract.featureset import Featureset
+import nkululeko.glob_conf as glob_conf
+
+
+class Whisper(Featureset):
+    """Class to extract whisper embeddings."""
+
+    def __init__(self, name, data_df, feat_type):
+        super().__init__(name, data_df, feat_type)
+        cuda = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = self.util.config_val("MODEL", "device", cuda)
+        self.model_initialized = False
+        if feat_type == "whisper":
+            self.feat_type = "whisper-base"
+        else:
+            self.feat_type = feat_type
+
+    def init_model(self):
+        # load model
+        self.util.debug("loading whisper model...")
+        model_name = f"openai/{self.feat_type}"
+        self.model = WhisperModel.from_pretrained(model_name).to(self.device)
+        print(f"intialized Whisper model on {self.device}")
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        self.model_initialized = True
+
+    def extract(self):
+        """Extract the features or load them from disk if present."""
+        store = self.util.get_path("store")
+        storage = f"{store}{self.name}.pkl"
+        extract = self.util.config_val("FEATS", "needs_feature_extraction", False)
+        no_reuse = eval(self.util.config_val("FEATS", "no_reuse", "False"))
+        if extract or no_reuse or not os.path.isfile(storage):
+            if not self.model_initialized:
+                self.init_model()
+            self.util.debug("extracting whisper embeddings, this might take a while...")
+            emb_series = []
+            for (file, start, end), _ in audeer.progress_bar(
+                self.data_df.iterrows(),
+                total=len(self.data_df),
+                desc=f"Running whisper on {len(self.data_df)} audiofiles",
+            ):
+                if end == pd.NaT:
+                    signal, sr = audiofile.read(file, offset=start)
+                else:
+                    signal, sr = audiofile.read(
+                        file, duration=end - start, offset=start
+                    )
+                emb = self.get_embeddings(signal, sr, file)
+                emb_series.append(emb)
+            # print(f"emb_series shape: {emb_series.shape}")
+            self.df = pd.DataFrame(emb_series, index=self.data_df.index)
+            # print(f"df shape: {self.df.shape}")
+            self.df.to_pickle(storage)
+            try:
+                glob_conf.config["DATA"]["needs_feature_extraction"] = "false"
+            except KeyError:
+                pass
+        else:
+            self.util.debug("reusing extracted wav2vec2 embeddings")
+            self.df = pd.read_pickle(storage)
+            if self.df.isnull().values.any():
+                nanrows = self.df.columns[self.df.isna().any()].tolist()
+                # print(nanrows)
+                self.util.error(
+                    f"got nan: {self.df.shape} {self.df.isnull().sum().sum()}"
+                )
+
+    def get_embeddings(self, signal, sampling_rate, file):
+        r"""Extract embeddings from raw audio signal."""
+        try:
+            with torch.no_grad():
+                embed_size = self.model.config.hidden_size
+                embed_columns = [f"whisper_{i}" for i in range(embed_size)]
+                inputs = self.feature_extractor(signal, sampling_rate=16000)[
+                    "input_features"
+                ][0]
+                inputs = torch.from_numpy(inputs).to(self.device).unsqueeze(0)
+                decoder_input_ids = (
+                    torch.tensor([[1, 1]]).to(self.device)
+                    * self.model.config.decoder_start_token_id
+                )
+                full_outputs = self.model(
+                    inputs,
+                    decoder_input_ids=decoder_input_ids,
+                    output_hidden_states=True,
+                )
+                outputs = full_outputs.encoder_last_hidden_state[0]
+                average_embeds = outputs.squeeze().mean(axis=0).cpu().detach().numpy()
+        except RuntimeError as re:
+            print(str(re))
+            self.util.error(f"couldn't extract file: {file}")
+        # print(f"y flattened shape: {y.ravel().shape}")
+        return average_embeds
+
+    def extract_sample(self, signal, sr):
+        self.init_model()
+        feats = self.get_embeddings(signal, sr, "no file")
+        return feats

nkululeko/feat_extract/featureset.py

@@ -7,13 +7,15 @@ import ast
 
 class Featureset:
     name = ""  # designation
-    df = None  # pandas dataframe to store the features (and indexed with the data from the sets)
+    df = None  # pandas dataframe to store the features
+    # (and indexed with the data from the sets)
     data_df = None  # dataframe to get audio paths
 
-    def __init__(self, name, data_df):
+    def __init__(self, name, data_df, feats_type):
         self.name = name
         self.data_df = data_df
         self.util = Util("featureset")
+        self.feats_types = feats_type
 
     def extract(self):
         pass
@@ -23,7 +25,8 @@ class Featureset:
         self.df = self.df[self.df.index.isin(self.data_df.index)]
         try:
             # use only some features
-            selected_features = ast.literal_eval(glob_conf.config["FEATS"]["features"])
+            selected_features = ast.literal_eval(
+                glob_conf.config["FEATS"]["features"])
             self.util.debug(f"selecting features: {selected_features}")
             sel_feats_df = pd.DataFrame()
             hit = False

nkululeko/feature_extractor.py

@@ -1,22 +1,25 @@
-"""
-feature_extractor.py
-
-Helper class to encapsulate feature extraction methods
+"""Extract acoustic features from audio samples.
 
+Extract acoustic features using several feature extractors 
+(appends the features column-wise)
 """
+
 import pandas as pd
 
 from nkululeko.utils.util import Util
 
 
 class FeatureExtractor:
-    """
-    Extract acoustic features from audio samples, using several feature extractors (appends the features column-wise)
+    """Extract acoustic features from audio samples.
+
+    Extract acoustic features using several feature extractors (appends the features column-wise).
+
     Args:
         data_df (pandas.DataFrame): dataframe with audiofile paths as index
-        feats_types (array of strings): designations of acoustic feature extractors to be used
-        data_name (string): names of databases that are extracted (for the caching)
-        feats_designation (string): the type of split (train/test), also is used for the cache name.
+        feats_types (List[str]): designations of acoustic feature extractors to be used
+        data_name (str): name of databases that are extracted (for caching)
+        feats_designation (str): the type of split (train/test), also is used for the cache name.
+
     Returns:
         df (pandas.DataFrame): dataframe with same index as data_df and acoustic features in columns
     """
@@ -25,7 +28,6 @@ class FeatureExtractor:
     df = None
     data_df = None  # dataframe to get audio paths
 
-    # def __init__
     def __init__(self, data_df, feats_types, data_name, feats_designation):
         self.data_df = data_df
         self.data_name = data_name
@@ -34,147 +36,80 @@ class FeatureExtractor:
         self.feats_designation = feats_designation
 
     def extract(self):
-        # feats_types = self.util.config_val_list('FEATS', 'type', ['os'])
-        self.featExtractor = None
         self.feats = pd.DataFrame()
-        _scale = True
         for feats_type in self.feats_types:
             store_name = f"{self.data_name}_{feats_type}"
-            if feats_type == "os":
-                from nkululeko.feat_extract.feats_opensmile import Opensmileset
-
-                self.featExtractor = Opensmileset(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "spectra":
-                from nkululeko.feat_extract.feats_spectra import Spectraloader
-
-                self.featExtractor = Spectraloader(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "trill":
-                from nkululeko.feat_extract.feats_trill import TRILLset
-
-                self.featExtractor = TRILLset(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type.startswith("wav2vec"):
-                from nkululeko.feat_extract.feats_wav2vec2 import Wav2vec2
-
-                self.featExtractor = Wav2vec2(
-                    f"{store_name}_{self.feats_designation}",
-                    self.data_df,
-                    feats_type,
-                )
-            elif feats_type.startswith("hubert"):
-                from nkululeko.feat_extract.feats_hubert import Hubert
-
-                self.featExtractor = Hubert(
-                    f"{store_name}_{self.feats_designation}",
-                    self.data_df,
-                    feats_type,
-                )
-
-            elif feats_type.startswith("wavlm"):
-                from nkululeko.feat_extract.feats_wavlm import Wavlm
-
-                self.featExtractor = Wavlm(
-                    f"{store_name}_{self.feats_designation}",
-                    self.data_df,
-                    feats_type,
-                )
-
-            elif feats_type.startswith("spkrec"):
-                from nkululeko.feat_extract.feats_spkrec import Spkrec
-
-                self.featExtractor = Spkrec(
-                    f"{store_name}_{self.feats_designation}",
-                    self.data_df,
-                    feats_type,
-                )
-            elif feats_type == "audmodel":
-                from nkululeko.feat_extract.feats_audmodel import AudModelSet
-
-                self.featExtractor = AudModelSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "auddim":
-                from nkululeko.feat_extract.feats_audmodel_dim import (
-                    AudModelDimSet,
-                )
-
-                self.featExtractor = AudModelDimSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "agender":
-                from nkululeko.feat_extract.feats_agender import (
-                    AudModelAgenderSet,
-                )
-
-                self.featExtractor = AudModelAgenderSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "agender_agender":
-                from nkululeko.feat_extract.feats_agender_agender import (
-                    AgenderAgenderSet,
-                )
-
-                self.featExtractor = AgenderAgenderSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "snr":
-                from nkululeko.feat_extract.feats_snr import SNRSet
-
-                self.featExtractor = SNRSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "mos":
-                from nkululeko.feat_extract.feats_mos import MOSSet
-
-                self.featExtractor = MOSSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "squim":
-                from nkululeko.feat_extract.feats_squim import SQUIMSet
-
-                self.featExtractor = SQUIMSet(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "clap":
-                from nkululeko.feat_extract.feats_clap import Clap
-
-                self.featExtractor = Clap(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "praat":
-                from nkululeko.feat_extract.feats_praat import Praatset
-
-                self.featExtractor = Praatset(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "mld":
-                from nkululeko.feat_extract.feats_mld import MLD_set
-
-                self.featExtractor = MLD_set(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            elif feats_type == "import":
-                from nkululeko.feat_extract.feats_import import Importset
-
-                self.featExtractor = Importset(
-                    f"{store_name}_{self.feats_designation}", self.data_df
-                )
-            else:
-                self.util.error(f"unknown feats_type: {feats_type}")
-
-            self.featExtractor.extract()
-            self.featExtractor.filter()
-            # remove samples that were not extracted by MLD
-            # self.df_test = self.df_test.loc[self.df_test.index.intersection(featExtractor_test.df.index)]
-            # self.df_train = self.df_train.loc[self.df_train.index.intersection(featExtractor_train.df.index)]
-            self.util.debug(f"{feats_type}: shape : {self.featExtractor.df.shape}")
-            self.feats = pd.concat([self.feats, self.featExtractor.df], axis=1)
+            self.feat_extractor = self._get_feat_extractor(store_name, feats_type)
+            self.feat_extractor.extract()
+            self.feat_extractor.filter()
+            self.feats = pd.concat([self.feats, self.feat_extractor.df], axis=1)
         return self.feats
 
     def extract_sample(self, signal, sr):
-        return self.featExtractor.extract_sample(signal, sr)
+        return self.feat_extractor.extract_sample(signal, sr)
+
+    def _get_feat_extractor(self, store_name, feats_type):
+        feat_extractor_class = self._get_feat_extractor_class(feats_type)
+        if feat_extractor_class is None:
+            self.util.error(f"unknown feats_type: {feats_type}")
+        return feat_extractor_class(
+            f"{store_name}_{self.feats_designation}", self.data_df, feats_type
+        )
+
+    def _get_feat_extractor_class(self, feats_type):
+        if feats_type == "os":
+            from nkululeko.feat_extract.feats_opensmile import Opensmileset
+
+            return Opensmileset
+
+        elif feats_type == "spectra":
+            from nkululeko.feat_extract.feats_spectra import Spectraloader
+
+            return Spectraloader
+
+        elif feats_type == "trill":
+            from nkululeko.feat_extract.feats_trill import TRILLset
+
+            return TRILLset
+
+        elif feats_type.startswith(
+            ("wav2vec2", "hubert", "wavlm", "spkrec", "whisper")
+        ):
+            return self._get_feat_extractor_by_prefix(feats_type)
+
+        elif feats_type == "xbow":
+            from nkululeko.feat_extract.feats_oxbow import Openxbow
+
+            return Openxbow
+
+        elif feats_type in (
+            "audmodel",
+            "auddim",
+            "agender",
+            "agender_agender",
+            "snr",
+            "mos",
+            "squim",
+            "clap",
+            "praat",
+            "mld",
+            "import",
+        ):
+            return self._get_feat_extractor_by_name(feats_type)
+        else:
+            return None
+
+    def _get_feat_extractor_by_prefix(self, feats_type):
+        prefix, _, ext = feats_type.partition("-")
+        from importlib import import_module
+
+        module = import_module(f"nkululeko.feat_extract.feats_{prefix.lower()}")
+        class_name = f"{prefix.capitalize()}"
+        return getattr(module, class_name)
+
+    def _get_feat_extractor_by_name(self, feats_type):
+        from importlib import import_module
+
+        module = import_module(f"nkululeko.feat_extract.feats_{feats_type.lower()}")
+        class_name = f"{feats_type.capitalize()}Set"
+        return getattr(module, class_name)