caul 0.1.0__tar.gz

caul-0.1.0/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.3
+Name: caul
+Version: 0.1.0
+Summary: Python implementation of an ASR service
+Author: Lion Summerbell
+Author-email: Lion Summerbell <lsummerbell@icij.org>
+Requires-Dist: nemo-toolkit[asr]
+Requires-Dist: torch>=2.10.0
+Requires-Dist: numpy
+Requires-Dist: torchaudio>=2.10.0
+Requires-Dist: torchcodec>=0.9.1
+Requires-Dist: librosa>=0.11.0
+Requires-Python: >=3.10.0, <3.14.0

caul-0.1.0/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "caul"
+version = "0.1.0"
+description = "Python implementation of an ASR service"
+requires-python = ">=3.10.0, <3.14.0"
+dependencies = [
+    "nemo-toolkit[asr]",
+    "torch>=2.10.0",
+    "numpy",
+    "torchaudio>=2.10.0",
+    "torchcodec>=0.9.1",
+    "librosa>=0.11.0",
+]
+authors = [
+    {name = "Lion Summerbell", email = "lsummerbell@icij.org"}
+]
+[build-system]
+requires = ["uv_build >= 0.9.22, <0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "black>=25.12.0",
+    "pre-commit>=4.5.1",
+    "pylint>=4.0.4",
+    "pytest>=9.0.2",
+]
+
+[tool.pylint."MESSAGES CONTROL"]
+disable = [
+    "missing-module-docstring"
+]
+
+[tool.pylint.main]
+fail-under = 8.89
+
+[tool.pylint.refactoring]
+max-nested-blocks = 5
+never-returning-functions = ["sys.exit", "argparse.parse_error"]
+
+[tool.pylint.similarities]
+# Docstrings are removed from the similarity computation
+ignore-docstrings = true
+# Comments are removed from the similarity computation
+ignore-comments = true
+
+[tool.black]
+target-version = ["py310"]

caul-0.1.0/src/caul/configs/__init__.py

@@ -0,0 +1,7 @@
+from caul.configs.parakeet import ParakeetConfig
+from caul.constant import PARAKEET
+
+
+MODEL_FAMILY_CONFIG_MAP = {
+    PARAKEET: ParakeetConfig(),
+}

caul-0.1.0/src/caul/configs/asr.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+
+from caul.constant import DEVICE_CPU
+
+if TYPE_CHECKING:
+    from caul.model_handlers.asr_model_handler import ASRModelHandler
+
+
+@dataclass
+class ASRConfig:
+    """Base config class"""
+
+    model_name: str
+    model_handler: "ASRModelHandler"
+    device: str | torch.device = DEVICE_CPU
+
+    def handler_from_config(self) -> "ASRModelHandler":
+        return (
+            self.model_handler(config=self) if self.model_handler is not None else None
+        )

caul-0.1.0/src/caul/configs/parakeet.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from caul.configs.asr import ASRConfig
+from caul.constant import EXPECTED_SAMPLE_RATE, PARAKEET_MODEL_REF
+from caul.model_handlers.parakeet import ParakeetModelHandler
+
+if TYPE_CHECKING:
+    from caul.model_handlers.asr_model_handler import ASRModelHandler
+
+
+@dataclass
+class ParakeetConfig(ASRConfig):
+
+    model_name: str = PARAKEET_MODEL_REF
+    model_handler: "ASRModelHandler" = ParakeetModelHandler
+    save_to_filesystem: bool = True
+    return_tensors: bool = True
+    sample_rate: int = EXPECTED_SAMPLE_RATE

caul-0.1.0/src/caul/constant.py

@@ -0,0 +1,31 @@
+# General
+
+DEVICE_CPU = "cpu"
+
+DEVICE_GPU = "gpu"
+
+DEVICE_MPS = "mps"
+
+EXPECTED_SAMPLE_RATE = 16000
+
+EXPECTED_FORMAT = "wav"
+
+EXPECTED_SAMPLE_MINUTE = EXPECTED_SAMPLE_RATE * 60
+
+# Parakeet
+
+PARAKEET_MODEL_REF = "nvidia/parakeet-tdt-0.6b-v3"
+
+PARAKEET_INFERENCE_MAX_DURATION_MIN = (
+    20  # actually 24, but we want to give ourselves some room
+)
+
+PARAKEET_INFERENCE_MAX_DURATION_KHZ = (
+    PARAKEET_INFERENCE_MAX_DURATION_MIN * EXPECTED_SAMPLE_MINUTE
+)
+
+PARAKEET = "parakeet"
+
+# Whisper
+
+WHISPER_CPP = "whisper-cpp"

caul-0.1.0/src/caul/exception.py

@@ -0,0 +1,6 @@
+class MissingModelSpecificationException(Exception):
+    """Raise if referencing a missing model"""
+
+
+class UnsupportedModelException(Exception):
+    """Raise if an unsupported model type is passed"""

caul-0.1.0/src/caul/filesystem.py

@@ -0,0 +1,33 @@
+from tempfile import mkstemp
+
+import torch
+import torchaudio
+
+from caul.constant import EXPECTED_SAMPLE_RATE, EXPECTED_FORMAT
+
+
+def save_tensor(audio_tensor: torch.Tensor) -> str:
+    """Filesystem routine for audio tensor; defaults to wav
+
+    :param audio_tensor: input tensor
+    :return: string file uri
+    """
+    # TODO: Change paths to run_id + tensor uuid + pagination
+    #  Allow for remote paths
+
+    _, file_path = mkstemp()
+
+    # torchcodec requires this
+    file_path = f"{file_path}.wav"
+
+    # Channel required as first dim
+    audio_tensor = audio_tensor.unsqueeze(0)
+
+    torchaudio.save(
+        file_path,
+        audio_tensor,
+        sample_rate=EXPECTED_SAMPLE_RATE,
+        format=EXPECTED_FORMAT,
+    )
+
+    return file_path

caul-0.1.0/src/caul/handler.py

@@ -0,0 +1,163 @@
+import logging
+
+from dataclasses import dataclass, field
+
+import torch
+
+import numpy as np
+
+from caul.configs.asr import ASRConfig
+from caul.exception import (
+    MissingModelSpecificationException,
+    UnsupportedModelException,
+)
+from caul.configs import MODEL_FAMILY_CONFIG_MAP
+from caul.tasks.inference.asr_inference import (
+    ASRModelHandlerResult,
+)
+from caul.model_handlers.asr_model_handler import ASRModelHandler, ASRModelHandlerResult
+from caul.utils import dict_key_fuzzy_match
+
+logger = logging.getLogger(__name__)
+
+
+class ASRHandler:
+    """ASRHandler class"""
+
+    # pylint: disable=R0913,R0917
+
+    def __init__(
+        self,
+        models: (
+            list[str | ASRModelHandler | ASRConfig] | str | ASRModelHandler | ASRConfig
+        ),
+        device: torch.device | str = None,
+        language_map: dict[str, int] = None,
+    ):
+        """Primary application handler class. Handles transcription agnostically.
+
+        :param models: Model_handler(s) or string reference(s)
+        :param device: cuda/cpu/mps
+        :param language_map: Map from ISO-639-3 language code to index of inference_handler
+        """
+        self.device = device
+
+        if language_map is None:
+            language_map = {}
+
+        self.language_map = language_map
+
+        self.model_handlers = []
+
+        if isinstance(models, list) and len(models) == 0:
+            raise MissingModelSpecificationException(
+                "At least one model name or model handler must be provided"
+            )
+
+        if not isinstance(models, list):
+            models = [models]
+
+        for model in models:
+            if isinstance(model, str):
+                supported_model_config = dict_key_fuzzy_match(
+                    MODEL_FAMILY_CONFIG_MAP, model
+                )
+
+                if supported_model_config is None:
+                    raise UnsupportedModelException(f"Unsupported model '{model}'")
+
+                # Set device after instantiation
+                supported_model_handler = supported_model_config.handler_from_config()
+                supported_model_handler.set_device(self.device)
+
+                self.model_handlers.append(supported_model_handler)
+            elif isinstance(model, ASRModelHandler):
+                self.model_handlers.append(model)
+            else:
+                raise UnsupportedModelException(f"Unsupported model type '{model}'")
+
+    def __repr__(self):
+        return f"<ASRHandler " f"models: {self.model_handlers} "
+
+    def startup(self):
+        """Run all model handler startup procedures"""
+        for model_handler in self.model_handlers:
+            model_handler.startup()
+
+    def shutdown(self):
+        """Garbage collect model handlers"""
+        self.model_handlers = []
+
+    def get_handler_by_language(self, language: str) -> ASRModelHandler:
+        """Get model_handler from language map or return first reference if language is not mapped
+
+        :param language: ISO-639-3 language code
+        :return: ASRModelHandler
+        """
+        reference_idx = self.language_map.get(
+            language, 0
+        )  # default to primary inference_handler when no language given
+
+        if len(self.model_handlers) <= reference_idx:
+            raise UnsupportedModelException(
+                "Language is mapped to a model index which does not exist"
+            )
+
+        return self.model_handlers[reference_idx]
+
+    def transcribe(
+        self,
+        inputs: list[np.ndarray | torch.Tensor | str] | np.ndarray | torch.Tensor | str,
+        languages: list[str] = None,
+    ) -> list[ASRModelHandlerResult]:
+        """Transcribe audio tensors or strings. Returns a tuple of (transcription, score). A list
+        of languages of len(inputs) may be passed to direct inputs to certain inference_handlers.
+
+        :param inputs: List of np.ndarray or torch.Tensor or str, or a singleton of same types
+        :param languages: List of ISO-639-3 language codes
+        :return: HandlerResult
+        """
+        if len(self.model_handlers) == 0:
+            raise MissingModelSpecificationException(
+                "At least one model name or model handler must be provided"
+            )
+
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+
+        audios_by_language = {}
+        model_handler_results_by_language = {}
+        batch_language_ordering = []
+        model_handler_results = []
+
+        if languages is None:
+            # Default to first model handler
+            return self.model_handlers[0].process(inputs)
+
+        # Sort by language where present, preserving original order for returning result
+        for idx, aud in enumerate(inputs):
+            language = languages[idx]
+
+            if language not in audios_by_language:
+                audios_by_language[language] = []
+
+            batch_language_ordering.append(language)
+            audios_by_language[language].append(aud)
+
+        # Run inference_handler on language batch
+        for language, audio_list in audios_by_language.items():
+            model_handler = self.get_handler_by_language(language)
+            model_handler_results_by_language[language] = model_handler.process(
+                audio_list
+            )
+
+        # For use with .pop()
+        batch_language_ordering.reverse()
+
+        # Reassemble and postprocess
+        for language in batch_language_ordering:
+            model_handler_result = model_handler_results_by_language[language].pop()
+
+            model_handler_results.append(model_handler_result)
+
+        return model_handler_results

caul-0.1.0/src/caul/model_handlers/asr_model_handler.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch
+
+from caul.configs.asr import ASRConfig
+from caul.model_handlers.helpers import ASRModelHandlerResult
+from caul.tasks.asr_task import ASRTask
+
+
+class ASRModelHandler(ABC):
+    """ASR model handler abstract"""
+
+    # pylint: disable=R0903
+
+    def __init__(self, config: "ASRConfig", *args, **kwargs):
+        self.config = config
+        self.tasks: list[ASRTask] = []
+
+    @abstractmethod
+    def startup(self):
+        """Generic method to load ASR resources"""
+
+    @abstractmethod
+    def shutdown(self):
+        """Generic method to unload ASR resources"""
+
+    def process(
+        self,
+        inputs: list[np.ndarray | torch.Tensor | str] | np.ndarray | torch.Tensor | str,
+    ) -> list[ASRModelHandlerResult]:
+        """Generic sequential processing method for ASR model handlers"""
+
+        output = inputs
+
+        for task in self.tasks:
+            output = task.process(output)
+
+        return output

caul-0.1.0/src/caul/model_handlers/helpers.py

@@ -0,0 +1,67 @@
+from dataclasses import dataclass
+
+from nemo.collections.asr.parts.utils import Hypothesis
+
+
+@dataclass
+class ASRModelHandlerResult:
+    """Base result class for ASR models"""
+
+    input_ordering: int = -1
+    transcription: list[tuple] = None
+    score: float = None
+
+
+@dataclass
+class ParakeetModelHandlerResult(ASRModelHandlerResult):
+    """Result handler for ParakeetInferenceHandler objects"""
+
+    def parse_parakeet_hypothesis(
+        self, hypothesis: Hypothesis
+    ) -> ASRModelHandlerResult:
+        """Parse a hypothesis returned by a Parakeet RNN model
+
+        :param hypothesis: Parakeet hypothesis
+        :return: copy of self
+        """
+        self.transcription = (
+            [
+                (s["start"], s["end"], s["segment"])
+                for s in hypothesis.timestamp.get("segment")
+            ]
+            if hypothesis.timestamp.get("segment") is not None
+            else [(0.0, 0.0, hypothesis.text)]
+        )
+        self.score = round(hypothesis.score, 2)
+
+        return self
+
+    def concat(self, model_result: ASRModelHandlerResult) -> ASRModelHandlerResult:
+        """Left fold with ParakeetModelHandlerResult object
+
+        :param model_result: ParakeetModelHandlerResult
+        :return: copy of self
+        """
+        if model_result is None:
+            return self
+
+        if self.transcription is None:
+            self.transcription = []
+
+        self.transcription += model_result.transcription
+
+        # We have to weight by total segment len
+        transcription_duration = self.transcription[-1][1]
+        model_result_duration = model_result.transcription[-1][1]
+        total_duration = transcription_duration + model_result_duration
+
+        self.score = round(
+            (
+                self.score * transcription_duration
+                + model_result.score * model_result_duration
+            )
+            / total_duration,
+            2,
+        )
+
+        return self

caul-0.1.0/src/caul/model_handlers/parakeet.py

@@ -0,0 +1,73 @@
+from typing import TYPE_CHECKING
+
+import torch
+
+from caul.constant import DEVICE_CPU, PARAKEET_MODEL_REF
+from caul.model_handlers.asr_model_handler import ASRModelHandler
+from caul.tasks.inference.parakeet_inference import ParakeetInferenceHandler
+from caul.tasks.postprocessing.parakeet_postprocessor import ParakeetPostprocessor
+from caul.tasks.preprocessing.parakeet_preprocessor import ParakeetPreprocessor
+
+if TYPE_CHECKING:
+    from caul.configs import ParakeetConfig
+
+
+class ParakeetModelHandler(ASRModelHandler):
+    """Model handler for Parakeet family"""
+
+    def __init__(
+        self,
+        config: "ParakeetConfig" = None,
+        model_name: str = PARAKEET_MODEL_REF,
+        device: str | torch.device = DEVICE_CPU,
+    ):
+        super().__init__(config=config)
+
+        if config is not None and config.model_name is not None:
+            model_name = config.model_name
+
+        self.model_name = model_name
+
+        if config is not None and config.device is not None:
+            device = config.device
+
+        if isinstance(device, str):
+            device = torch.device(device)
+
+        self.device = device
+
+        self.preprocessor = ParakeetPreprocessor(
+            save_to_filesystem=config.save_to_filesystem,
+            return_tensors=config.return_tensors,
+        )
+        self.inference_handler = ParakeetInferenceHandler(
+            model_name=config.model_name, device=config.device
+        )
+        self.postprocessor = ParakeetPostprocessor()
+
+        self.tasks = [self.preprocessor, self.inference_handler, self.postprocessor]
+
+    def set_device(self, device: str | torch.device = DEVICE_CPU):
+        """Set/change device here and on inference_handler
+
+        :param device: device to use
+        """
+        if isinstance(device, str):
+            device = torch.device(device)
+
+        self.device = device
+
+        self.inference_handler.set_device(device)
+
+        return self
+
+    def startup(self):
+        """Load model"""
+        self.inference_handler.load()
+
+    def shutdown(self):
+        """Shut down"""
+        self.preprocessor = None
+        self.inference_handler = None
+        self.postprocessor = None
+        self.tasks = []

caul-0.1.0/src/caul/tasks/asr_task.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class ASRTask(ABC):
+    """Generic ASR task"""
+
+    # pylint: disable=R0903
+
+    @abstractmethod
+    def process(self, inputs: Any, *args, **kwargs) -> list:
+        """Generic processing task"""

caul-0.1.0/src/caul/tasks/inference/asr_inference.py

@@ -0,0 +1,24 @@
+from abc import abstractmethod
+
+from caul.model_handlers.asr_model_handler import ASRModelHandlerResult
+from caul.tasks.asr_task import ASRTask
+
+
+class ASRInferenceHandler(ASRTask):
+    """Abstract for ASR inference"""
+
+    @abstractmethod
+    def process(self, inputs: list, *args, **kwargs) -> list[ASRModelHandlerResult]:
+        """
+
+        :param inputs: List of inference inputs
+        :return: ASRModelHandlerResult
+        """
+
+    @abstractmethod
+    def load(self):
+        """Load model"""
+
+    @abstractmethod
+    def unload(self):
+        """Unload model"""

caul-0.1.0/src/caul/tasks/inference/parakeet_inference.py

@@ -0,0 +1,83 @@
+import torch
+
+import nemo.collections.asr as nemo_asr
+
+from caul.constant import DEVICE_CPU
+from caul.model_handlers.helpers import ParakeetModelHandlerResult
+from caul.tasks.inference.asr_inference import (
+    ASRInferenceHandler,
+)
+from caul.tasks.preprocessing.helpers import PreprocessedInput
+
+
+class ParakeetInferenceHandler(ASRInferenceHandler):
+    """Inference handler for NVIDIA's Parakeet family of ASR models. Supports up to 24 minutes of
+    audio (batched or unbatched) in a single pass. Assumes that audio inputs (wav files or tensors)
+    are single-channel with a sample rate of 16000—this last is very important for segmenting.
+    """
+
+    def __init__(self, model_name: str, device: str | torch.device = DEVICE_CPU):
+        self.model_name = model_name
+
+        if isinstance(device, str):
+            device = torch.device(device)
+
+        self.device = device
+        self.model = None
+
+    def load(self):
+        """Load model; default to CPU where no device is present"""
+        device = self.device
+
+        if device is None:
+            device = DEVICE_CPU
+
+        self.model = nemo_asr.models.ASRModel.from_pretrained(
+            self.model_name, map_location=torch.device(device)
+        ).eval()
+
+    def unload(self):
+        """Unload model"""
+        self.model = None
+
+    def set_device(self, device: str | torch.device = DEVICE_CPU):
+        """Set/change device"""
+        if isinstance(device, str):
+            device = torch.device(device)
+
+        self.device = device
+
+        return self
+
+    def process(
+        self,
+        inputs: list[list[PreprocessedInput]] | list[PreprocessedInput],
+        timestamps: bool = True,
+    ) -> list[ParakeetModelHandlerResult]:
+        """Transcribe a batch of audio tensors or file names of max duration <= 20 minutes
+
+        :param inputs: List of np.ndarray or torch.Tensor or str, or singleton of same types
+        :param timestamps: Whether to include timestamps with transcriptions
+        :return: List of results
+        """
+        if len(inputs) == 0:
+            return []
+
+        if isinstance(inputs[0], PreprocessedInput):
+            inputs = [inputs]
+
+        transcriptions = []
+
+        for input_batch in inputs:
+            hypotheses = self.model.transcribe(
+                [i.tensor.to(self.device) for i in input_batch], timestamps=timestamps
+            )
+            # Get timestamped segments if available, otherwise default to whole text
+            for idx, hyp in enumerate(hypotheses):
+                input_ordering_idx = input_batch[idx].metadata.input_ordering
+                model_result = ParakeetModelHandlerResult(
+                    input_ordering=input_ordering_idx
+                ).parse_parakeet_hypothesis(hyp)
+                transcriptions.append(model_result)
+
+        return transcriptions

caul-0.1.0/src/caul/tasks/inference/whisper_cpp_inference.py

@@ -0,0 +1,20 @@
+from caul.tasks.inference.asr_inference import (
+    ASRInferenceHandler,
+    ASRModelHandlerResult,
+)
+
+
+class WhisperCPPInferenceHandler(ASRInferenceHandler):
+    """Handler for WhisperCPP; wrapper round subprocess calls"""
+
+    # pylint: disable=R0903
+
+    def process(
+        self,
+        inputs: list[str],
+    ) -> list[ASRModelHandlerResult]:
+        """List of np.ndarray or torch.Tensor or str, or a singleton of same types
+
+        :param inputs: List of np.ndarray or torch.Tensor or str, or a singleton of same types
+        :return:
+        """

caul-0.1.0/src/caul/tasks/postprocessing/parakeet_postprocessor.py

@@ -0,0 +1,55 @@
+from functools import reduce
+from itertools import groupby
+
+from caul.tasks.asr_task import ASRTask
+from caul.tasks.inference.parakeet_inference import ParakeetModelHandlerResult
+
+
+class ParakeetPostprocessor(ASRTask):
+    """Postprocessing logic for ParakeetInferenceHandler output"""
+
+    def process(
+        self, inputs: list[ParakeetModelHandlerResult]
+    ) -> list[ParakeetModelHandlerResult]:
+        """Process indexed ParakeetInferenceHandler results and return them in their original
+        ordering
+
+        :param inputs: List of parakeet model results
+        :return: list of parakeet model results in input ordering
+        """
+
+        return self.map_results_to_inputs(inputs)
+
+    @staticmethod
+    def map_results_to_inputs(
+        batched_results: list[ParakeetModelHandlerResult],
+    ) -> list[ParakeetModelHandlerResult]:
+        """Remap unordered and segmented tensors to original inputs for return
+
+        :param batched_results: list of unordered ParakeetModelHandlerResult, still
+        segmented
+        :return: list[ParakeetModelHandlerResult]
+        """
+        unbatched_results = []
+
+        # Sort in order before batching
+        batched_results = sorted(batched_results, key=lambda r: r.input_ordering)
+
+        # Concat segmented tensors
+        results_grouped_by_index = groupby(
+            batched_results, key=lambda r: r.input_ordering
+        )
+
+        for _, group_results in results_grouped_by_index:
+            group_results = list(group_results)
+
+            merged_results = (
+                reduce(lambda l, r: l.concat(r), group_results)
+                if len(group_results) > 1
+                else group_results[0]
+            )
+            unbatched_results.append(merged_results)
+
+        # TODO: Drop index from result
+
+        return unbatched_results

caul-0.1.0/src/caul/tasks/preprocessing/helpers.py

@@ -0,0 +1,32 @@
+import datetime
+import uuid
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+
+@dataclass
+class InputMetadata:
+    """Preprocessed input metadata"""
+
+    input_ordering: int
+    duration: int
+    start_time: int = 0
+    end_time: int = 0
+    preprocessed_at: str = (
+        datetime.datetime.now(datetime.timezone.utc).replace(microsecond=0).isoformat()
+    )
+    uuid: str = uuid.uuid4().hex
+    input_format: str = None
+    input_file_path: str = None
+    preprocessed_file_path: str = None
+
+
+@dataclass
+class PreprocessedInput:
+    """Preprocessed input wrapper"""
+
+    metadata: InputMetadata
+    tensor: Optional[torch.Tensor | list] = None

caul-0.1.0/src/caul/tasks/preprocessing/parakeet_preprocessor.py

@@ -0,0 +1,258 @@
+import librosa
+import torch
+
+import numpy as np
+import torchaudio
+
+from caul.constant import (
+    PARAKEET_INFERENCE_MAX_DURATION_KHZ,
+    EXPECTED_SAMPLE_MINUTE,
+    EXPECTED_SAMPLE_RATE,
+    PARAKEET_INFERENCE_MAX_DURATION_MIN,
+)
+from caul.filesystem import save_tensor
+from caul.tasks.asr_task import ASRTask
+from caul.tasks.preprocessing.helpers import PreprocessedInput, InputMetadata
+
+
+class ParakeetPreprocessor(ASRTask):
+    """Preprocessing logic for ParakeetInferenceHandler inputs"""
+
+    def __init__(
+        self,
+        save_to_filesystem: bool = True,
+        return_tensors: bool = True,
+        sample_rate: int = EXPECTED_SAMPLE_RATE,
+    ):
+        super().__init__()
+
+        self.save_to_filesystem = save_to_filesystem
+        self.return_tensors = return_tensors
+        self.sample_rate = sample_rate
+
+    def process(
+        self,
+        inputs: list[np.ndarray | torch.Tensor | str] | np.ndarray | torch.Tensor | str,
+        input_sample_rates: list[int] | int = None,
+    ) -> list[list[PreprocessedInput]]:
+        """Segment and batch audio inputs
+
+        :param inputs: List of np.ndarray or torch.Tensor or str, or singleton of same types
+        :param input_sample_rates: sample rate(s) of audio inputs
+        :return: batches of indexed preprocessed audio tensors (input_idx, preprocessed_input)
+        """
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+
+        preprocessed_inputs = self.preprocess_inputs(inputs, input_sample_rates)
+        batches = self.batch_audio_tensors(preprocessed_inputs)
+
+        return batches
+
+    def preprocess_inputs(
+        self,
+        inputs: list[np.ndarray | torch.Tensor | str],
+        input_sample_rates: list[int] = None,
+    ) -> list[PreprocessedInput]:
+        """Accepts audio inputs as a list of file paths, np.ndarray, or torch.Tensor, converting to
+        torch.Tensor, normalizing, segmenting inputs longer than 20 minutes (just under Parakeet's
+        max) first by silences or with overlaps where not available, and batching segments
+
+        :param inputs: List of np.ndarray or torch.Tensor or str, or a singleton of same types
+        :param input_sample_rates: sample rate(s) of audio inputs
+        :return: List of processed inputs
+        """
+        preprocessed_inputs = []
+
+        # Load arrays and divide into max_length segments
+        for input_idx, audio_input in enumerate(inputs):
+            input_file_path = None
+            new_file_path = None
+            input_format = None
+
+            # Load audio files as arrays
+            if isinstance(audio_input, str):
+                input_file_path = audio_input
+                input_format = (
+                    input_file_path.split(".")[-1]
+                    if len(input_file_path.split(".")) > 1
+                    else None
+                )
+                audio_input, sample_rate = torchaudio.load(audio_input)
+
+            if isinstance(audio_input, np.ndarray):
+                audio_input = torch.Tensor(audio_input)
+
+            # Normalize
+            if input_sample_rates is not None and len(input_sample_rates) > input_idx:
+                sample_rate = input_sample_rates[input_idx]
+            else:
+                sample_rate = self.sample_rate
+
+            audio_input = self.normalize(audio_input, sample_rate)
+
+            # Segment where necessary
+            duration_khz = audio_input.shape[-1]
+            tensor_segments = [audio_input]
+
+            if duration_khz > PARAKEET_INFERENCE_MAX_DURATION_KHZ:
+                tensor_segments = self.segment_audio_tensor(audio_input)
+
+            for tensor_segment in tensor_segments:
+                # Create temporary filesystem reference if applicable
+                if self.save_to_filesystem:
+                    new_file_path = save_tensor(tensor_segment)
+
+                if not self.return_tensors:
+                    tensor_segment = None
+
+                # Create preprocessed input
+                metadata = InputMetadata(
+                    input_ordering=input_idx,
+                    duration=duration_khz / EXPECTED_SAMPLE_MINUTE,
+                    input_format=input_format,
+                    input_file_path=input_file_path,
+                    preprocessed_file_path=new_file_path,
+                )
+
+                preprocessed_input = PreprocessedInput(
+                    tensor=tensor_segment,
+                    metadata=metadata,
+                )
+
+                preprocessed_inputs.append(preprocessed_input)
+
+        return preprocessed_inputs
+
+    def normalize(self, audio_tensor: torch.Tensor, sample_rate: int) -> torch.Tensor:
+        """Normalize audio_tensor (single channel, sample rate = 16000)
+
+        :param audio_tensor: input tensor
+        :param sample_rate: input sample rate
+        :return: normalized tensor
+        """
+        if sample_rate != self.sample_rate:
+            audio_tensor = self.resample_waveform(audio_tensor, sample_rate)
+
+        # Stereo dims (channels, aud_length); need mono (aud_length)
+        if len(audio_tensor.shape) > 1:
+            audio_tensor = audio_tensor.squeeze(0)
+
+        return audio_tensor
+
+    @staticmethod
+    def segment_audio_tensor(
+        audio_tensor: torch.Tensor,
+        frame_len: int = 2048,
+        silence_thresh_db: int = 35,
+        hop_len: int = 512,
+        kept_silence_len_secs: int = 0.15,
+        min_silence_len_secs: int = 0.5,
+        max_segment_len_secs: int = EXPECTED_SAMPLE_MINUTE
+        * PARAKEET_INFERENCE_MAX_DURATION_MIN,
+    ) -> list[torch.Tensor]:
+        """Splits on silences with librosa, falling back to overlaps where min segments
+        are not sufficient to safely divide audio.
+
+        :param audio_tensor: input tensor
+        :param frame_len: number of samples per analysis frame
+        :param silence_thresh_db: max decibel value
+        :param hop_len: number of samples between analysis frames
+        :param kept_silence_len_secs: number of seconds to keep silence
+        :param min_silence_len_secs: minimum seconds to keep silence
+        :param max_segment_len_secs: maximum seconds to keep silence
+        :return: list of tensor segments
+        """
+        # TODO: Implement fallback to overlaps
+        tensor_segments = []
+
+        # Intervals between silences
+        nonsilent_intervals = librosa.effects.split(
+            audio_tensor.numpy(),
+            top_db=silence_thresh_db,
+            frame_length=frame_len,
+            hop_length=hop_len,
+        )
+
+        merged = []
+        min_silence_sample_len = int(min_silence_len_secs * EXPECTED_SAMPLE_MINUTE)
+        kept_silence_sample_len = int(kept_silence_len_secs * EXPECTED_SAMPLE_MINUTE)
+        max_segment_sample_len = int(max_segment_len_secs * EXPECTED_SAMPLE_MINUTE)
+
+        # Merge intervals separated by short silences
+        for start, end in nonsilent_intervals:
+            if len(merged) == 0:
+                merged.append((start, end))
+            else:
+                _, prev_end = merged[-1]
+                if start - prev_end < min_silence_sample_len:
+                    merged[-1][1] = end
+                else:
+                    merged.append((start, end))
+
+        # Segment controlling max length
+        for start, end in merged:
+            start = max(0, start - kept_silence_sample_len)
+            end = min(audio_tensor.shape[-1], end + kept_silence_sample_len)
+
+            while end - start > max_segment_sample_len:
+                segment_end = start + max_segment_sample_len
+                tensor_segment = audio_tensor[start:segment_end]
+
+                tensor_segments.append(tensor_segment)
+
+        return tensor_segments
+
+    @staticmethod
+    def batch_audio_tensors(  # pylint: disable=R0914
+        preprocessed_inputs: list[PreprocessedInput],
+    ) -> list[list[PreprocessedInput]]:
+        """Batch audio tensors by duration, 20 minutes max per batch, optimizing for tightly packed
+        batches.
+
+        :param preprocessed_inputs: list of PreprocessedInput
+        :return: list of list[PreprocessedInput]
+        """
+
+        # Sort by duration
+        preprocessed_inputs = sorted(
+            preprocessed_inputs, key=lambda p: p.metadata.duration, reverse=True
+        )
+
+        # Now this becomes a bin-packing minimization problem. We'll use a variant of best-fit
+        # decreasing.
+
+        bins = [[]]
+        bins_len = [0]
+
+        # With each pass, choose a bin by maximizing remaining space
+        for preprocessed_input in preprocessed_inputs:
+            bin_len_diffs = []
+
+            for bin_len in bins_len:
+                bin_len_diffs.append(PARAKEET_INFERENCE_MAX_DURATION_MIN - bin_len)
+
+            if max(bin_len_diffs) <= preprocessed_input.metadata.duration:
+                bins.append([])
+                bins_len.append(0)
+                bin_len_diffs.append(PARAKEET_INFERENCE_MAX_DURATION_MIN)
+
+            max_diff_idx = np.argmax(bin_len_diffs)
+
+            bins[max_diff_idx].append(preprocessed_input)
+
+            bins_len[max_diff_idx] += preprocessed_input.metadata.duration
+
+        return bins
+
+    def resample_waveform(
+        self, waveform: torch.Tensor, sample_rate: int
+    ) -> torch.Tensor:
+        """Resample when sample rate is not 16000
+
+        :param waveform: torch.Tensor
+        :param sample_rate: int
+        :return: resampled torch.Tensor
+        """
+        transform = torchaudio.transforms.Resample(sample_rate, self.sample_rate)
+        return transform(waveform)

caul-0.1.0/src/caul/utils.py

@@ -0,0 +1,24 @@
+from typing import Any
+
+
+def dict_key_fuzzy_match(dict_obj: dict, search_key: str) -> Any | None:
+    """Match a dict key name fuzzily (returning first match)
+
+    :param dict_obj: dictionary to search in
+    :param search_key: key name
+    :return: key value (if exists)
+    """
+
+    if search_key in dict_obj:
+        return dict_obj[search_key]
+
+    fuzzy_matches = [
+        dict_value
+        for dict_key, dict_value in dict_obj.items()
+        if dict_key in search_key or search_key in dict_key
+    ]
+
+    if len(fuzzy_matches) > 0:
+        return fuzzy_matches[0]
+
+    return None