lfeats 0.1.3__tar.gz → 0.2.0__tar.gz

{lfeats-0.1.3 → lfeats-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lfeats
-Version: 0.1.3
+Version: 0.2.0
 Summary: A unified interface to extract hidden representations from speech foundation models
 Project-URL: Documentation, https://takenori-y.github.io/lfeats/stable/
 Project-URL: Source, https://github.com/takenori-y/lfeats
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.10
+Requires-Dist: filelock>=3.10.0
 Requires-Dist: huggingface-hub>=0.23.0
 Requires-Dist: hydra-core>=1.3.0
 Requires-Dist: hyperpyyaml>=0.0.1

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/cli.py

@@ -7,8 +7,10 @@
 
 import argparse
 import logging
+import mimetypes
 import os
 import sys
+from pathlib import Path
 
 logger = logging.getLogger("lfeats")
 
@@ -21,14 +23,23 @@ def get_arguments() -> argparse.Namespace:
     parser.add_argument(
         "source",
         type=str,
-        help="The source audio file, directory, or .scp file containing file paths.",
+        help="The source audio file, directory, or scp file containing file paths.",
     )
     parser.add_argument(
         "--output_dir",
         type=str,
-        default=None,
+        default=".",
         help="The directory where the extracted features will be saved.",
     )
+    parser.add_argument(
+        "--subdir_offset",
+        type=int,
+        default=None,
+        help=(
+            "The offset from the beginning of the input file path to start creating "
+            "subdirectories in the output directory."
+        ),
+    )
     parser.add_argument(
         "--output_format",
         type=str,
@@ -127,6 +138,9 @@ def main() -> None:
     """Perform the main feature extraction process."""
     args = get_arguments()
 
+    if args.subdir_offset is not None and args.subdir_offset < 0:
+        raise ValueError("Subdir offset must be a non-negative integer.")
+
     logging.basicConfig(
         level=logging.ERROR if args.quiet else logging.INFO,
         format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
@@ -141,11 +155,12 @@ def main() -> None:
 
     # Get the list of input files from the source argument.
     if os.path.isfile(args.source):
-        if args.source.endswith(".scp"):
+        mime_type, _ = mimetypes.guess_type(args.source)
+        if mime_type is not None and mime_type.startswith("audio/"):
+            input_files = [args.source]
+        else:
             with open(args.source) as f:
                 input_files = [line.strip() for line in f if line.strip()]
-        else:
-            input_files = [args.source]
     elif os.path.isdir(args.source):
         input_files = []
         for root, _, files in os.walk(args.source):
@@ -168,10 +183,6 @@ def main() -> None:
     else:
         raise ValueError(f"Invalid layers argument: {args.layers}")
 
-    # Prepare the output directory if specified.
-    if args.output_dir is not None:
-        os.makedirs(args.output_dir, exist_ok=True)
-
     output_ext = {
         "npz": "npz",
         "pt": "pt",
@@ -199,10 +210,27 @@ def main() -> None:
             num_errors += 1
             continue
 
+        if args.subdir_offset is None:
+            subdir = ""
+        else:
+            path = Path(input_file).parent
+            # Remove the root part of the path.
+            dirs = path.relative_to(path.anchor).parts
+            if args.subdir_offset >= len(dirs):
+                logger.error(
+                    f"Subdir offset {args.subdir_offset} is too large for file: "
+                    f"{input_file}. Skipping."
+                )
+                num_errors += 1
+                continue
+            subdirs = dirs[args.subdir_offset :]
+            subdir = Path(*subdirs)
+
+        output_dir = os.path.join(args.output_dir, subdir)
+        os.makedirs(output_dir, exist_ok=True)
+
         base, _ = os.path.splitext(os.path.basename(input_file))
-        output_file = f"{base}.{output_ext}"
-        if args.output_dir is not None:
-            output_file = os.path.join(args.output_dir, output_file)
+        output_file = os.path.join(output_dir, f"{base}.{output_ext}")
 
         try:
             audio, sample_rate = load_audio(input_file)

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/interfaces/extractor.py

@@ -90,6 +90,18 @@ class Extractor:
         """
         self.model_manager.get_model().load(self.cache_dir, quiet)
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.model_manager.to(device)
+        self.resampler_manager.to(device)
+
     def __call__(
         self,
         source: np.ndarray | torch.Tensor | Audio,

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/interfaces/resampler.py

@@ -45,6 +45,17 @@ class Resampler:
                 f"Supported resamplers are: {[k for k in RESAMPLER_MAP.keys()]}"
             ) from e
 
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.resampler_manager.to(device)
+
     def __call__(
         self,
         source: np.ndarray | torch.Tensor | Audio,

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/base.py

@@ -25,6 +25,7 @@ class BaseModel(ABC):
         """
         self.device = device
 
+        self.model = None
         self._model_id = None  # To be defined in subclasses
 
     @abstractmethod
@@ -42,6 +43,24 @@ class BaseModel(ABC):
         """
         raise NotImplementedError
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        if self.model is not None and hasattr(self.model, "to"):
+            self.model.to(device)
+            if hasattr(self.model, "device"):
+                try:
+                    self.model.device = device
+                except AttributeError:
+                    pass
+
     def extract_features(self, audio: Audio, layers: list[int]) -> Features:
         """Extract features from the input audio data.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/contentvec.py

@@ -68,8 +68,6 @@ class ContentVecModel(FrameLevelFeatureModel):
         )
         self._model_id = f"contentvec-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/data2vec.py

@@ -58,9 +58,6 @@ class Data2VecModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Data2VecVariant, Data2VecVariant.BASE)
         self._model_id = f"data2vec-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/data2vec2.py

@@ -58,9 +58,6 @@ class Data2Vec2Model(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Data2Vec2Variant, Data2Vec2Variant.BASE)
         self._model_id = f"data2vec2-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/ecapa_tdnn.py

@@ -40,8 +40,6 @@ class EcapaTDNNModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, EcapaTDNNVariant, EcapaTDNNVariant.BASE)
         self._model_id = f"ecapa-tdnn-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class EcapaTDNNModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-ecapa-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             self.model.eval()
-            self.model.to(self.device)
 
     def extract_features_impl(self, audio: Audio, layers: list[int]) -> Features:
         """Extract features from the input audio using the model.

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/emotion2vec.py

@@ -55,8 +55,6 @@ class Emotion2VecModel(FrameLevelFeatureModel):
         )
         self._model_id = f"emotion2vec-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,7 +76,7 @@ class Emotion2VecModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename=filename,
                 repo_type="model",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         setup_third_party_path()

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/emotion2vec_plus.py

@@ -58,8 +58,6 @@ class Emotion2VecPlusModel(FrameLevelFeatureModel):
         )
         self._model_id = f"emotion2vec+-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -81,13 +79,13 @@ class Emotion2VecPlusModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename="model.pt",
                 repo_type="model",
-                local_dir=os.path.join(model_dir, sanitize(self.model_id)),
+                cache_dir=os.path.join(model_dir, sanitize(self.model_id)),
             )
             config = hf_hub_download(
                 repo_id=repo_id,
                 filename="config.yaml",
                 repo_type="model",
-                local_dir=os.path.join(model_dir, sanitize(self.model_id)),
+                cache_dir=os.path.join(model_dir, sanitize(self.model_id)),
             )
 
         from hyperpyyaml import load_hyperpyyaml

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/hubert.py

@@ -60,8 +60,6 @@ class HuBERTModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, HuBERTVariant, HuBERTVariant.BASE)
         self._model_id = f"hubert-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/manager.py

@@ -32,6 +32,19 @@ class ModelManager:
 
         self._cache: dict[str, BaseModel] = {}
 
+    def to(self, device: str) -> None:
+        """Move all models to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the models to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        for model in self._cache.values():
+            model.to(device)
+
     def get_model(self) -> BaseModel:
         """Get the model instance.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/next_tdnn.py

@@ -66,8 +66,6 @@ class NeXtTDNNModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, NeXtTDNNVariant, NeXtTDNNVariant.BASE_V2)
         self._model_id = f"next-tdnn-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/r_spin.py

@@ -63,8 +63,6 @@ class RSpinModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, RSpinVariant, RSpinVariant.WAVLM_256)
         self._model_id = f"r-spin-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/r_vector.py

@@ -40,8 +40,6 @@ class RVectorModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, RVectorVariant, RVectorVariant.BASE)
         self._model_id = f"r-vector-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class RVectorModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-resnet-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             self.model.eval()
-            self.model.to(self.device)
 
     def extract_features_impl(self, audio: Audio, layers: list[int]) -> Features:
         """Extract features from the input audio using the model.

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/spidr.py

@@ -4,12 +4,11 @@
 """A module for the SpidR model."""
 
 from enum import Enum
-from typing import Any
 
 import torch
 
 from ..interfaces.types import Audio, Features
-from ..utils.io import set_torch_hub_dir
+from ..utils.io import safe_torch_hub_load
 from ..utils.validation import validate_enum
 from .base import FrameLevelFeatureModel
 
@@ -40,8 +39,6 @@ class SpidRModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, SpidRVariant, SpidRVariant.BASE)
         self._model_id = f"spidr-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -57,10 +54,9 @@ class SpidRModel(FrameLevelFeatureModel):
         if self.model is not None:
             return
 
-        with set_torch_hub_dir(model_dir):
-            self.model: Any = torch.hub.load(
-                "facebookresearch/spidr", "spidr_base", verbose=not quiet
-            )
+        self.model = safe_torch_hub_load(
+            "facebookresearch/spidr", "spidr_base", model_dir, quiet=quiet
+        )
         self.model.eval()
         self.model.to(self.device)
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/spin.py

@@ -117,8 +117,6 @@ class SpinModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, SpinVariant, SpinVariant.HUBERT_256)
         self._model_id = f"spin-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -139,7 +137,7 @@ class SpinModel(FrameLevelFeatureModel):
                 repo_id="vectominist/spin_ckpt",
                 filename=self.variant.checkpoint_filename,
                 repo_type="dataset",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         with lightning_mock_context():
@@ -147,11 +145,12 @@ class SpinModel(FrameLevelFeatureModel):
                 model_path, map_location=torch.device("cpu"), weights_only=False
             )
 
-        from lfeats.third_party.s3prl.util.download import set_dir
+        from lfeats.third_party.s3prl.util.download import set_dir, set_progress
         from lfeats.third_party.spin.model import SpinModel as _SpinModel
         from lfeats.third_party.spin.util import len_to_padding
 
         set_dir(model_dir)
+        set_progress(not quiet)
         self.model = _SpinModel(checkpoint["hyper_parameters"])
         self.model.load_state_dict(checkpoint["state_dict"])
         self.model.eval()

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/sslzip.py

@@ -63,6 +63,19 @@ class SSLZipModel(FrameLevelFeatureModel):
         self.upstream = HuBERTModel(variant=HuBERTVariant.BASE.value, device=device)
         self.model = None
 
+    def to(self, device: str) -> None:
+        """Move the model to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the model to (e.g., 'cpu' or 'cuda'). Note that the ONNX
+            model will be loaded on the specified device when calling the 'load' method.
+
+        """
+        self.device = device
+        self.upstream.to(device)
+
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -86,7 +99,7 @@ class SSLZipModel(FrameLevelFeatureModel):
                 repo_id=repo_id,
                 filename=filename,
                 repo_type="model",
-                local_dir=model_dir,
+                cache_dir=model_dir,
             )
 
         import onnxruntime as ort

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/unispeech_sat.py

@@ -56,8 +56,6 @@ class UniSpeechSATModel(FrameLevelFeatureModel):
         )
         self._model_id = f"unispeech-sat-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/wav2vec2.py

@@ -64,9 +64,6 @@ class Wav2Vec2Model(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, Wav2Vec2Variant, Wav2Vec2Variant.BASE)
         self._model_id = f"wav2vec2-{self.variant.value}"
 
-        self.processor = None
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/wavlm.py

@@ -54,8 +54,6 @@ class WavLMModel(FrameLevelFeatureModel):
         self.variant = validate_enum(variant, WavLMVariant, WavLMVariant.BASE_PLUS)
         self._model_id = f"wavlm-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/whisper.py

@@ -58,7 +58,6 @@ class WhisperModel(FrameLevelFeatureModel):
         self._model_id = f"whisper-{self.variant.value}"
 
         self.processor = None
-        self.model = None
 
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/models/x_vector.py

@@ -40,8 +40,6 @@ class XVectorModel(UtteranceLevelFeatureModel):
         self.variant = validate_enum(variant, XVectorVariant, XVectorVariant.BASE)
         self._model_id = f"x-vector-{self.variant.value}"
 
-        self.model = None
-
     def load(self, model_dir: str, quiet: bool = False) -> None:
         """Load the model from the specified directory.
 
@@ -78,11 +76,11 @@ class XVectorModel(UtteranceLevelFeatureModel):
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-xvect-voxceleb",
                 fetch_config=fetch_config,
+                run_opts={"device": self.device},
             )
             if self.model is None:
                 raise RuntimeError("Failed to load the model.")
             self.model.eval()
-            self.model.to(self.device)
 
     def extract_features_impl(self, audio: Audio, layers: list[int]) -> Features:
         """Extract features from the input audio using the model.

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/resamplers/base.py

@@ -35,6 +35,17 @@ class BaseResampler(ABC):
         self.dst_rate = dst_rate
         self.device = device
 
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+
     def resample(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/resamplers/lilfilter.py

@@ -56,7 +56,19 @@ class LilFilterResampler(BaseResampler):
             dst_rate,
             dtype=torch.float32,
         )
-        self.resampler.weights = self.resampler.weights.to(self.device)
+        self.to(device)
+
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        super().to(device)
+        self.resampler.weights = self.resampler.weights.to(device)
 
     def resample_impl(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/resamplers/manager.py

@@ -32,6 +32,19 @@ class ResamplerManager:
 
         self._cache: dict[tuple[int, int], BaseResampler] = {}
 
+    def to(self, device: str) -> None:
+        """Move all resamplers to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resamplers to (e.g., 'cpu' or 'cuda').
+
+        """
+        self.device = device
+        for resampler in self._cache.values():
+            resampler.to(device)
+
     def get_resampler(self, src_rate: int, dst_rate: int) -> BaseResampler:
         """Get the resampler instance.
 

{lfeats-0.1.3 → lfeats-0.2.0}/lfeats/resamplers/torchaudio.py

@@ -73,7 +73,20 @@ class TorchAudioResampler(BaseResampler):
 
         self.resampler = torchaudio.transforms.Resample(
             orig_freq=src_rate, new_freq=dst_rate, dtype=torch.float32, **params
-        ).to(device)
+        )
+        self.to(device)
+
+    def to(self, device: str) -> None:
+        """Move the resampler to the specified device.
+
+        Parameters
+        ----------
+        device : str
+            The device to move the resampler to (e.g., 'cpu' or 'cuda').
+
+        """
+        super().to(device)
+        self.resampler.to(device)
 
     def resample_impl(self, audio: Audio) -> Audio:
         """Resample the given audio to the target sample rate.