nkululeko 0.94.2__py3-none-any.whl → 0.95.0__py3-none-any.whl

nkululeko/augmenting/resampler.py

@@ -17,7 +17,7 @@ class Resampler:
     def __init__(self, df, replace, not_testing=True):
         self.SAMPLING_RATE = 16000
         self.df = df
-        self.util = Util("resampler", has_config=not_testing)
+        self.util = Util("resampler", has_config=not not_testing)
         self.util.warn(f"all files might be resampled to {self.SAMPLING_RATE}")
         self.not_testing = not_testing
         self.replace = (
@@ -30,7 +30,7 @@ class Resampler:
         files = self.df.index.get_level_values(0).values
         # replace = eval(self.util.config_val("RESAMPLE", "replace", "False"))
         replace = self.replace
-        if self.not_testing:
+        if not self.not_testing:
             store = self.util.get_path("store")
         else:
             store = "./"
@@ -67,17 +67,28 @@ class Resampler:
             self.df = self.df.set_index(
                 self.df.index.set_levels(new_files, level="file")
             )
-            target_file = self.util.config_val("RESAMPLE", "target", "resampled.csv")
-            # remove encoded labels
-            target = self.util.config_val("DATA", "target", "emotion")
-            if "class_label" in self.df.columns:
-                self.df = self.df.drop(columns=[target])
-                self.df = self.df.rename(columns={"class_label": target})
-            # save file
-            self.df.to_csv(target_file)
-            self.util.debug(
-                "saved resampled list of files to" f" {os.path.abspath(target_file)}"
-            )
+            if not self.not_testing:
+                target_file = self.util.config_val(
+                    "RESAMPLE", "target", "resampled.csv"
+                )
+                # remove encoded labels
+                target = self.util.config_val("DATA", "target", "emotion")
+                if "class_label" in self.df.columns:
+                    self.df = self.df.drop(columns=[target])
+                    self.df = self.df.rename(columns={"class_label": target})
+                # save file
+                self.df.to_csv(target_file)
+                self.util.debug(
+                    "saved resampled list of files to"
+                    f" {os.path.abspath(target_file)}"
+                )
+            else:
+                # When running from command line, save to simple resampled.csv
+                target_file = "resampled.csv"
+                self.df.to_csv(target_file)
+                self.util.debug(
+                    f"saved resampled list of files to {os.path.abspath(target_file)}"
+                )
         self.util.debug(f"resampled {succes} files, {error} errors")
 
 
@@ -91,7 +102,7 @@ def main():
         df_sample.index, allow_nat=False
     )
     df_sample.head(10)
-    resampler = Resampler(df_sample, not_testing=False)
+    resampler = Resampler(df_sample, False, not_testing=False)
     resampler.resample()
     shutil.copyfile(testfile, "tmp.resample_result.wav")
     shutil.copyfile("tmp.wav", testfile)

nkululeko/autopredict/ap_emotion.py

@@ -0,0 +1,36 @@
+"""
+A predictor for emotion classification.
+Uses emotion2vec models for emotion prediction.
+"""
+
+import ast
+
+import nkululeko.glob_conf as glob_conf
+from nkululeko.feature_extractor import FeatureExtractor
+from nkululeko.utils.util import Util
+
+
+class EmotionPredictor:
+    """
+    EmotionPredictor
+    predicting emotion with emotion2vec models
+    """
+
+    def __init__(self, df):
+        self.df = df
+        self.util = Util("emotionPredictor")
+
+    def predict(self, split_selection):
+        self.util.debug(f"predicting emotion for {split_selection} samples")
+        feats_name = "_".join(ast.literal_eval(glob_conf.config["DATA"]["databases"]))
+        
+        self.feature_extractor = FeatureExtractor(
+            self.df, ["emotion2vec-large"], feats_name, split_selection
+        )
+        emotion_df = self.feature_extractor.extract()
+        
+        pred_emotion = ["neutral"] * len(emotion_df)
+        
+        return_df = self.df.copy()
+        return_df["emotion_pred"] = pred_emotion
+        return return_df

nkululeko/autopredict/ap_text.py

@@ -0,0 +1,45 @@
+"""A predictor for text.
+
+Currently based on whisper model.
+"""
+
+import ast
+
+import torch
+
+from nkululeko.feature_extractor import FeatureExtractor
+import nkululeko.glob_conf as glob_conf
+from nkululeko.utils.util import Util
+
+
+class TextPredictor:
+    """TextPredictor.
+
+    predicting text with the whisper model
+    """
+
+    def __init__(self, df, util=None):
+        self.df = df
+        if util is not None:
+            self.util = util
+        else:
+            # create a new util instance
+            # this is needed to access the config and other utilities
+            # in the autopredict module
+            self.util = Util("textPredictor")
+        from nkululeko.autopredict.whisper_transcriber import Transcriber
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = self.util.config_val("MODEL", "device", device)
+        self.transcriber = Transcriber(
+            device=device,
+            language=self.util.config_val("EXP", "language", "en"),
+            util=self.util,
+        )
+    def predict(self, split_selection):
+        self.util.debug(f"predicting text for {split_selection} samples")
+        df = self.transcriber.transcribe_index(
+            self.df.index
+        )
+        return_df = self.df.copy()
+        return_df["text"] = df["text"].values
+        return return_df

nkululeko/autopredict/whisper_transcriber.py

@@ -0,0 +1,81 @@
+import os
+
+import pandas as pd
+import torch
+from tqdm import tqdm
+import whisper
+
+import audeer
+import audiofile
+
+from nkululeko.utils.util import Util
+
+
+class Transcriber:
+    def __init__(self, model_name="turbo", device=None, language="en", util=None):
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = whisper.load_model(model_name, device=device)
+        self.language = language
+        self.util = util
+
+    def transcribe_file(self, audio_path):
+        """Transcribe the audio file at the given path.
+
+        :param audio_path: Path to the audio file to transcribe.
+        :return: Transcription text.
+        """
+        result = self.model.transcribe(
+            audio_path, language=self.language, without_timestamps=True)
+        result = result["text"].strip()
+        return result
+
+    def transcribe_array(self, signal, sampling_rate):
+        """Transcribe the audio file at the given path.
+
+        :param audio_path: Path to the audio file to transcribe.
+        :return: Transcription text.
+        """
+        tmporary_path = "temp.wav"
+        audiofile.write(
+            "temp.wav", signal, sampling_rate, format="wav")
+        result = self.transcribe_file(tmporary_path)
+        return result
+
+    def transcribe_index(self, index:pd.Index) ->  pd.DataFrame:
+        """Transcribe the audio files in the given index.
+
+        :param index: Index containing tuples of (file, start, end).
+        :return: DataFrame with transcriptions indexed by the original index.
+        :rtype: pd.DataFrame
+        """
+        file_name = ""
+        seg_index = 0
+        transcriptions = []
+        transcriber_cache = audeer.mkdir(
+            audeer.path(self.util.get_path("cache"), "transcriptions"))
+        for idx, (file, start, end) in enumerate(
+            tqdm(index.to_list())
+        ):
+            if file != file_name:
+                file_name = file
+                seg_index = 0
+            cache_name = audeer.basename_wo_ext(file)+str(seg_index)
+            cache_path = audeer.path(transcriber_cache, cache_name + ".json")
+            if os.path.isfile(cache_path):
+                transcription = self.util.read_json(cache_path)["transcription"]
+            else:
+                dur = end.total_seconds() - start.total_seconds()
+                y, sr = audiofile.read(file, offset=start, duration=dur)
+                transcription = self.transcribe_array(
+                    y, sr)
+                self.util.save_json(cache_path, 
+                                {"transcription": transcription, 
+                                 "file": file, 
+                                 "start": start.total_seconds(), 
+                                 "end": end.total_seconds()})
+            transcriptions.append(transcription)
+            seg_index += 1
+
+        df = pd.DataFrame({"text":transcriptions}, index=index)
+        return df

nkululeko/constants.py

@@ -1,2 +1,2 @@
-VERSION="0.94.2"
+VERSION="0.95.0"
 SAMPLING_RATE = 16000

nkululeko/experiment.py

@@ -513,7 +513,7 @@ class Experiment:
 
     def autopredict(self):
         """Predict labels for samples with existing models and add to the dataframe."""
-        sample_selection = self.util.config_val("PREDICT", "split", "all")
+        sample_selection = self.util.config_val("PREDICT", "sample_selection", "all")
         if sample_selection == "all":
             df = pd.concat([self.df_train, self.df_test])
         elif sample_selection == "train":
@@ -569,6 +569,11 @@ class Experiment:
 
                 predictor = STOIPredictor(df)
                 df = predictor.predict(sample_selection)
+            elif target == "text":
+                from nkululeko.autopredict.ap_text import TextPredictor
+
+                predictor = TextPredictor(df, self.util)
+                df = predictor.predict(sample_selection)
             elif target == "arousal":
                 from nkululeko.autopredict.ap_arousal import ArousalPredictor
 
@@ -584,6 +589,11 @@ class Experiment:
 
                 predictor = DominancePredictor(df)
                 df = predictor.predict(sample_selection)
+            elif target == "emotion":
+                from nkululeko.autopredict.ap_emotion import EmotionPredictor
+
+                predictor = EmotionPredictor(df)
+                df = predictor.predict(sample_selection)
             else:
                 self.util.error(f"unknown auto predict target: {target}")
         return df
@@ -668,11 +678,27 @@ class Experiment:
 
         # check if a scatterplot should be done
         scatter_var = eval(self.util.config_val("EXPL", "scatter", "False"))
+
+        # Priority: use [EXPL][scatter.target] if available, otherwise use [DATA][target] value
+        if hasattr(self, "target") and self.target != "none":
+            default_scatter_target = f"['{self.target}']"
+        else:
+            default_scatter_target = "['class_label']"
+
         scatter_target = self.util.config_val(
-            "EXPL", "scatter.target", "['class_label']"
+            "EXPL", "scatter.target", default_scatter_target
         )
+
+        if scatter_target == default_scatter_target:
+            self.util.debug(
+                f"scatter.target using default from [DATA][target]: {scatter_target}"
+            )
+        else:
+            self.util.debug(
+                f"scatter.target from [EXPL][scatter.target]: {scatter_target}"
+            )
         if scatter_var:
-            scatters = ast.literal_eval(glob_conf.config["EXPL"]["scatter"])
+            scatters = ast.literal_eval(scatter_target)
             scat_targets = ast.literal_eval(scatter_target)
             plots = Plots()
             for scat_target in scat_targets:
@@ -692,6 +718,30 @@ class Experiment:
                             df_feats, df_labels, f"{scat_target}_bins", scatter
                         )
 
+        # check if t-SNE plot should be generated
+        tsne = eval(self.util.config_val("EXPL", "tsne", "False"))
+        if tsne:
+            target_column = self.util.config_val("DATA", "target", "emotion")
+            plots = Plots()
+            self.util.debug("generating t-SNE plot...")
+            plots.scatter_plot(df_feats, df_labels, target_column, "tsne")
+
+        # check if UMAP plot should be generated
+        umap_plot = eval(self.util.config_val("EXPL", "umap", "False"))
+        if umap_plot:
+            target_column = self.util.config_val("DATA", "target", "emotion")
+            plots = Plots()
+            self.util.debug("generating UMAP plot...")
+            plots.scatter_plot(df_feats, df_labels, target_column, "umap")
+
+        # check if PCA plot should be generated
+        pca_plot = eval(self.util.config_val("EXPL", "pca", "False"))
+        if pca_plot:
+            target_column = self.util.config_val("DATA", "target", "emotion")
+            plots = Plots()
+            self.util.debug("generating PCA plot...")
+            plots.scatter_plot(df_feats, df_labels, target_column, "pca")
+
     def _check_scale(self):
         self.util.save_to_store(self.feats_train, "feats_train")
         self.util.save_to_store(self.feats_test, "feats_test")

nkululeko/explore.py

@@ -8,6 +8,8 @@ The script supports the following configuration options:
 - `no_warnings`: If set to `True`, it will ignore all warnings during the exploration.
 - `feature_distributions`: If set to `True`, it will generate plots of the feature distributions.
 - `tsne`: If set to `True`, it will generate a t-SNE plot of the feature space.
+- `umap`: If set to `True`, it will generate a UMAP plot of the feature space.
+- `pca`: If set to `True`, it will generate a PCA plot of the feature space.
 - `scatter`: If set to `True`, it will generate a scatter plot of the feature space.
 - `spotlight`: If set to `True`, it will generate a 'spotlight' plot of the feature space.
 - `shap`: If set to `True`, it will generate SHAP feature importance plots.
@@ -59,10 +61,12 @@ def main():
 
         warnings.filterwarnings("ignore")
     needs_feats = False
+    experiment_loaded = False
     try:
         # load the experiment
         expr.load(f"{util.get_save_name()}")
         needs_feats = True
+        experiment_loaded = True
     except FileNotFoundError:
         # first time: load the data
         expr.load_datasets()
@@ -73,20 +77,35 @@ def main():
             f"train shape : {expr.df_train.shape}, test shape:{expr.df_test.shape}"
         )
 
-        plot_feats = eval(util.config_val("EXPL", "feature_distributions", "False"))
-        tsne = eval(util.config_val("EXPL", "tsne", "False"))
-        scatter = eval(util.config_val("EXPL", "scatter", "False"))
-        shap = eval(util.config_val("EXPL", "shap", "False"))
-        model_type = util.config_val("EXPL", "model", False)
-        plot_tree = eval(util.config_val("EXPL", "plot_tree", "False"))
-        needs_feats = False
-        if plot_feats or tsne or scatter or model_type or plot_tree or shap:
-            # these investigations need features to explore
+    # Check exploration settings regardless of whether experiment was loaded or not
+    plot_feats = eval(util.config_val("EXPL", "feature_distributions", "False"))
+    tsne_plot = eval(util.config_val("EXPL", "tsne", "False"))
+    umap_plot = eval(util.config_val("EXPL", "umap", "False"))
+    pca_plot = eval(util.config_val("EXPL", "pca", "False"))
+    scatter = eval(util.config_val("EXPL", "scatter", "False"))
+    shap = eval(util.config_val("EXPL", "shap", "False"))
+    model_type = util.config_val("EXPL", "model", False)
+    plot_tree = eval(util.config_val("EXPL", "plot_tree", "False"))
+
+    if (
+        plot_feats
+        or tsne_plot
+        or umap_plot
+        or pca_plot
+        or scatter
+        or model_type
+        or plot_tree
+        or shap
+    ):
+        # these investigations need features to explore
+        if not experiment_loaded or not needs_feats:
             expr.extract_feats()
-            needs_feats = True
-            # explore
-            # expr.init_runmanager()
-            # expr.runmgr.do_runs()
+        needs_feats = True
+        # explore
+        if shap:
+            # SHAP analysis requires a trained model
+            expr.init_runmanager()
+            expr.runmgr.do_runs()
     expr.analyse_features(needs_feats)
     expr.store_report()
     print("DONE")

nkululeko/feat_extract/feats_analyser.py

@@ -1,5 +1,6 @@
 # feats_analyser.py
 import ast
+import os
 
 import matplotlib.pyplot as plt
 import pandas as pd
@@ -76,17 +77,37 @@ class FeatureAnalyser:
             self.util.to_pickle(shap_values, name)
         else:
             shap_values = self.util.from_pickle(name)
-        # plt.figure()
-        plt.close("all")
-        plt.tight_layout()
-        shap.plots.bar(shap_values)
-        fig_dir = self.util.get_path("fig_dir") + "../"  # one up because of the runs
-        exp_name = self.util.get_exp_name(only_data=True)
+        # Create SHAP summary plot instead
+        fig, ax = plt.subplots(figsize=(10, 6))
+        shap.plots.bar(shap_values, ax=ax, show=False)
+        fig_dir = os.path.join(self.util.get_path("fig_dir"), "..")
+
         format = self.util.config_val("PLOT", "format", "png")
-        filename = f"_SHAP_{model.name}"
-        filename = f"{fig_dir}{exp_name}{filename}.{format}"
-        plt.savefig(filename)
-        plt.close()
+        feat_type = self.util.get_feattype_name()
+        filename = f"SHAP_{feat_type}_{model.name}.{format}"
+        filename = os.path.join(fig_dir, filename)
+
+        fig.savefig(filename, dpi=300, bbox_inches="tight")
+        plt.close(fig)
+
+        # print and save SHAP feature importance
+        max_feat_num = len(self.features.columns)
+        shap_importance_values = shap_values.abs.mean(0).values
+
+        feature_cols = self.features.columns
+        feature_importance = pd.DataFrame(
+            shap_importance_values[:max_feat_num],
+            index=feature_cols,
+            columns=["importance"],
+        ).sort_values("importance", ascending=False)
+
+        self.util.debug(
+            f"SHAP analysis, features = {feature_importance.index.tolist()}"
+        )
+        # Save to CSV (save all features, not just top ones)
+        csv_filename = os.path.join(fig_dir, f"SHAP_{feat_type}_importance_{model.name}.csv")
+        feature_importance.to_csv(csv_filename)
+        self.util.debug(f"Saved SHAP feature importance to {csv_filename}")
         self.util.debug(f"plotted SHAP feature importance to {filename}")
 
     def analyse(self):
@@ -120,6 +141,12 @@ class FeatureAnalyser:
                     covariance_type = self.util.config_val(
                         "MODEL", "GMM_covariance_type", "full"
                     )
+                    allowed_cov_types = ["full", "tied", "diag", "spherical"]
+                    if covariance_type not in allowed_cov_types:
+                        self.util.error(
+                            f"Invalid covariance_type '{covariance_type}', must be one of {allowed_cov_types}. Using default 'full'."
+                        )
+                        covariance_type = "full"
                     model = mixture.GaussianMixture(
                         n_components=n_components, covariance_type=covariance_type
                     )
@@ -156,7 +183,7 @@ class FeatureAnalyser:
                     from sklearn.svm import SVC
 
                     c = float(self.util.config_val("MODEL", "C_val", "1.0"))
-                    model = SVC(kernel="linear", C=c, gamma="scale")
+                    model = SVC(kernel="linear", C=c, gamma="scale", random_state=42)
                     result_importances[model_s] = self._get_importance(
                         model, permutation
                     )
@@ -165,7 +192,7 @@ class FeatureAnalyser:
                         plots = Plots()
                         plots.plot_tree(model, self.features)
                 elif model_s == "tree":
-                    model = DecisionTreeClassifier()
+                    model = DecisionTreeClassifier(random_state=42)
                     result_importances[model_s] = self._get_importance(
                         model, permutation
                     )
@@ -176,7 +203,9 @@ class FeatureAnalyser:
                 elif model_s == "xgb":
                     from xgboost import XGBClassifier
 
-                    model = XGBClassifier(enable_categorical=True, tree_method="hist")
+                    model = XGBClassifier(
+                        enable_categorical=True, tree_method="hist", random_state=42
+                    )
                     self.labels = self.labels.astype("category")
                     result_importances[model_s] = self._get_importance(
                         model, permutation
@@ -263,13 +292,12 @@ class FeatureAnalyser:
             title += "\n based on feature permutation"
         ax.set(title=title)
         plt.tight_layout()
-        fig_dir = self.util.get_path("fig_dir") + "../"  # one up because of the runs
-        exp_name = self.util.get_exp_name(only_data=True)
+        fig_dir = self.util.get_path("fig_dir")
         format = self.util.config_val("PLOT", "format", "png")
-        filename = f"_EXPL_{model_name}"
+        filename = f"EXPL_{model_name}"
         if permutation:
             filename += "_perm"
-        filename = f"{fig_dir}{exp_name}{filename}.{format}"
+        filename = f"{fig_dir}{filename}.{format}"
         plt.savefig(filename)
         fig = ax.figure
         fig.clear()

nkululeko/feat_extract/feats_emotion2vec.py

@@ -3,7 +3,6 @@
 # choices for feat_type = "emotion2vec", "emotion2vec-large", "emotion2vec-base", "emotion2vec-seed"
 
 # requirements:
-# pip install "modelscope>=1.9.5,<2.0.0"
 # pip install funasr
 
 import os
@@ -43,27 +42,30 @@ class Emotion2vec(Featureset):
         except ImportError:
             self.util.error(
                 "FunASR is required for emotion2vec features. "
-                "Please install with: pip install funasr modelscope"
+                "Please install with: pip install funasr"
             )
 
-        # Map feat_type to model names
+        # Map feat_type to model names on HuggingFace
         model_mapping = {
-            "emotion2vec": "iic/emotion2vec_base",
-            "emotion2vec-base": "iic/emotion2vec_base_finetuned",
-            "emotion2vec-seed": "iic/emotion2vec_plus_seed",
-            "emotion2vec-large": "iic/emotion2vec_plus_large",
+            "emotion2vec": "emotion2vec/emotion2vec_base",
+            "emotion2vec-base": "emotion2vec/emotion2vec_base",
+            "emotion2vec-seed": "emotion2vec/emotion2vec_plus_seed",
+            "emotion2vec-large": "emotion2vec/emotion2vec_plus_large",
         }
 
         # Get model path from config or use default mapping
         model_path = self.util.config_val(
             "FEATS",
             "emotion2vec.model",
-            model_mapping.get(self.feat_type, "iic/emotion2vec_base"),
+            model_mapping.get(self.feat_type, "emotion2vec/emotion2vec_base"),
         )
 
         try:
-            # Initialize the FunASR model for emotion2vec
-            self.model = AutoModel(model=model_path)
+            # Initialize the FunASR model for emotion2vec using HuggingFace Hub
+            self.model = AutoModel(
+                model=model_path,
+                hub="hf"  # Use HuggingFace Hub instead of ModelScope
+            )
             self.util.debug(f"initialized emotion2vec model: {model_path}")
             self.model_initialized = True
         except Exception as e:
@@ -131,7 +133,9 @@ class Emotion2vec(Featureset):
                 import tempfile
                 import soundfile as sf
 
-                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+                with tempfile.NamedTemporaryFile(
+                    suffix=".wav", delete=False
+                ) as tmp_file:
                     sf.write(tmp_file.name, signal_np, sampling_rate)
                     audio_path = tmp_file.name
             else:
@@ -152,11 +156,20 @@ class Emotion2vec(Featureset):
                             embeddings = np.array(embeddings)
                         return embeddings.flatten()
                     else:
-                        # Fallback to create default embedding
-                        return np.array([0.0] * 768)
+                        # Fallback based on model type
+                        if 'large' in self.feat_type.lower():
+                            return np.array([0.0] * 1024)
+                        else:
+                            return np.array([0.0] * 768)
                 else:
-                    self.util.error(f"No result from emotion2vec model for file: {file}")
-                    return np.array([0.0] * 768)
+                    self.util.error(
+                        f"No result from emotion2vec model for file: {file}"
+                    )
+                    # Fallback based on model type
+                    if 'large' in self.feat_type.lower():
+                        return np.array([0.0] * 1024)
+                    else:
+                        return np.array([0.0] * 768)
 
             finally:
                 # Clean up temporary file if we created one
@@ -166,36 +179,40 @@ class Emotion2vec(Featureset):
         except Exception as e:
             print(f"Error processing {file}: {str(e)}")
             self.util.error(f"couldn't extract file: {file}, error: {str(e)}")
-            return np.array([0.0] * 768)
+            # Return appropriate dimension based on model type
+            if 'large' in self.feat_type.lower():
+                return np.array([0.0] * 1024)
+            else:
+                return np.array([0.0] * 768)
 
     def extract_sample(self, signal, sr):
         """Extract features from a single sample."""
         if not self.model_initialized:
             self.init_model()
-        
+
         # Save signal as temporary file for emotion2vec
         import tempfile
         import soundfile as sf
-        
+
         try:
             # Convert tensor to numpy if needed
             if torch.is_tensor(signal):
                 signal_np = signal.squeeze().numpy()
             else:
                 signal_np = signal.squeeze()
-                
+
             # Handle multi-channel audio
             if signal_np.ndim > 1:
                 signal_np = signal_np[0]
-                
+
             with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
                 sf.write(tmp_file.name, signal_np, sr)
-                
+
                 # Extract using the emotion2vec model
                 res = self.model.generate(
                     tmp_file.name, granularity="utterance", extract_embedding=True
                 )
-                
+
                 # Get embeddings from result
                 if isinstance(res, list) and len(res) > 0:
                     embeddings = res[0].get("feats", None)
@@ -203,12 +220,20 @@ class Emotion2vec(Featureset):
                         if isinstance(embeddings, list):
                             embeddings = np.array(embeddings)
                         return embeddings.flatten()
-                
-                return np.array([0.0] * 768)
-                
+
+                # Fallback based on model type
+                if 'large' in self.feat_type.lower():
+                    return np.array([0.0] * 1024)
+                else:
+                    return np.array([0.0] * 768)
+
         except Exception as e:
             print(f"Error in extract_sample: {str(e)}")
-            return np.array([0.0] * 768)
+            # Return appropriate dimension based on model type
+            if 'large' in self.feat_type.lower():
+                return np.array([0.0] * 1024)
+            else:
+                return np.array([0.0] * 768)
         finally:
             # Clean up temporary file
             if tmp_file is not None:  # Check if tmp_file was created