lattifai 0.4.6__py3-none-any.whl → 1.0.0__py3-none-any.whl

lattifai/__init__.py

@@ -1,28 +1,40 @@
 import sys
 import warnings
+from importlib.metadata import version
 
+# Re-export I/O classes
+from .caption import Caption
+
+# Re-export client classes
+from .client import LattifAI
+
+# Re-export config classes
+from .config import (
+    AUDIO_FORMATS,
+    MEDIA_FORMATS,
+    VIDEO_FORMATS,
+    AlignmentConfig,
+    CaptionConfig,
+    ClientConfig,
+    DiarizationConfig,
+    MediaConfig,
+)
 from .errors import (
     AlignmentError,
     APIError,
     AudioFormatError,
     AudioLoadError,
     AudioProcessingError,
+    CaptionParseError,
+    CaptionProcessingError,
     ConfigurationError,
     DependencyError,
     LatticeDecodingError,
     LatticeEncodingError,
     LattifAIError,
     ModelLoadError,
-    SubtitleParseError,
-    SubtitleProcessingError,
 )
-from .io import SubtitleIO
-
-try:
-    from importlib.metadata import version
-except ImportError:
-    # Python < 3.8
-    from importlib_metadata import version
+from .logging import get_logger, set_log_level, setup_logger
 
 try:
     __version__ = version("lattifai")
@@ -54,28 +66,25 @@ def _check_and_install_k2():
 _check_and_install_k2()
 
 
-# Lazy import for LattifAI to avoid dependency issues during basic import
-def __getattr__(name):
-    if name == "LattifAI":
-        from .client import LattifAI
-
-        return LattifAI
-    if name == "AsyncLattifAI":
-        from .client import AsyncLattifAI
-
-        return AsyncLattifAI
-    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
-
-
 __all__ = [
-    "LattifAI",  # noqa: F822
-    "AsyncLattifAI",  # noqa: F822
+    # Client classes
+    "LattifAI",
+    # Config classes
+    "AlignmentConfig",
+    "ClientConfig",
+    "CaptionConfig",
+    "DiarizationConfig",
+    "MediaConfig",
+    "AUDIO_FORMATS",
+    "VIDEO_FORMATS",
+    "MEDIA_FORMATS",
+    # Error classes
     "LattifAIError",
     "AudioProcessingError",
     "AudioLoadError",
     "AudioFormatError",
-    "SubtitleProcessingError",
-    "SubtitleParseError",
+    "CaptionProcessingError",
+    "CaptionParseError",
     "AlignmentError",
     "LatticeEncodingError",
     "LatticeDecodingError",
@@ -83,6 +92,12 @@ __all__ = [
     "DependencyError",
     "APIError",
     "ConfigurationError",
-    "SubtitleIO",
+    # Logging
+    "setup_logger",
+    "get_logger",
+    "set_log_level",
+    # I/O
+    "Caption",
+    # Version
     "__version__",
 ]

lattifai/alignment/__init__.py

@@ -0,0 +1,6 @@
+"""Alignment module for LattifAI forced alignment."""
+
+from .lattice1_aligner import Lattice1Aligner
+from .segmenter import Segmenter
+
+__all__ = ["Lattice1Aligner", "Segmenter"]

lattifai/alignment/lattice1_aligner.py

@@ -0,0 +1,119 @@
+"""Lattice-1 Aligner implementation."""
+
+from typing import Any, List, Optional, Tuple
+
+import colorful
+import torch
+
+from lattifai.audio2 import AudioData
+from lattifai.caption import Supervision
+from lattifai.config import AlignmentConfig
+from lattifai.errors import (
+    AlignmentError,
+    LatticeDecodingError,
+    LatticeEncodingError,
+)
+from lattifai.utils import _resolve_model_path
+
+from .lattice1_worker import _load_worker
+from .tokenizer import _load_tokenizer
+
+ClientType = Any
+
+
+class Lattice1Aligner(object):
+    """Synchronous LattifAI client with config-driven architecture."""
+
+    def __init__(
+        self,
+        config: AlignmentConfig,
+    ) -> None:
+        self.config = config
+
+        if config.client_wrapper is None:
+            raise ValueError("AlignmentConfig.client_wrapper is not set. It must be initialized by the client.")
+
+        client_wrapper = config.client_wrapper
+        model_path = _resolve_model_path(config.model_name)
+
+        self.tokenizer = _load_tokenizer(client_wrapper, model_path, config.model_name, config.device)
+        self.worker = _load_worker(model_path, config.device)
+
+        self.frame_shift = self.worker.frame_shift
+
+    def emission(self, audio: torch.Tensor) -> torch.Tensor:
+        return self.worker.emission(audio.to(self.worker.device))
+
+    def alignment(
+        self,
+        audio: AudioData,
+        supervisions: List[Supervision],
+        split_sentence: Optional[bool] = False,
+        return_details: Optional[bool] = False,
+        emission: Optional[torch.Tensor] = None,
+        offset: float = 0.0,
+        verbose: bool = True,
+    ) -> Tuple[List[Supervision], List[Supervision]]:
+        """
+        Perform alignment on audio and supervisions.
+
+        Args:
+            audio: Audio file path
+            supervisions: List of supervision segments to align
+            split_sentence: Enable sentence re-splitting
+
+        Returns:
+            Tuple of (supervisions, alignments)
+
+        Raises:
+            LatticeEncodingError: If lattice graph generation fails
+            AlignmentError: If audio alignment fails
+            LatticeDecodingError: If lattice decoding fails
+        """
+        try:
+            if verbose:
+                print(colorful.cyan("🔗 Step 2: Creating lattice graph from segments"))
+            try:
+                supervisions, lattice_id, lattice_graph = self.tokenizer.tokenize(
+                    supervisions, split_sentence=split_sentence
+                )
+                if verbose:
+                    print(colorful.green(f"         ✓ Generated lattice graph with ID: {lattice_id}"))
+            except Exception as e:
+                text_content = " ".join([sup.text for sup in supervisions]) if supervisions else ""
+                raise LatticeEncodingError(text_content, original_error=e)
+
+            if verbose:
+                print(colorful.cyan(f"🔍 Step 3: Searching lattice graph with media: {audio}"))
+            try:
+                lattice_results = self.worker.alignment(audio, lattice_graph, emission=emission, offset=offset)
+                if verbose:
+                    print(colorful.green("         ✓ Lattice search completed"))
+            except Exception as e:
+                raise AlignmentError(
+                    f"Audio alignment failed for {audio}",
+                    media_path=str(audio),
+                    context={"original_error": str(e)},
+                )
+
+            if verbose:
+                print(colorful.cyan("🎯 Step 4: Decoding lattice results to aligned segments"))
+            try:
+                alignments = self.tokenizer.detokenize(
+                    lattice_id, lattice_results, supervisions=supervisions, return_details=return_details
+                )
+                if verbose:
+                    print(colorful.green(f"         ✓ Successfully aligned {len(alignments)} segments"))
+            except LatticeDecodingError as e:
+                print(colorful.red("         x Failed to decode lattice alignment results"))
+                raise e
+            except Exception as e:
+                print(colorful.red("         x Failed to decode lattice alignment results"))
+                raise LatticeDecodingError(lattice_id, original_error=e)
+
+            return (supervisions, alignments)
+
+        except (LatticeEncodingError, AlignmentError, LatticeDecodingError):
+            raise
+        except Exception as e:
+            raise e

lattifai/{workers/lattice1_alpha.py → alignment/lattice1_worker.py}

@@ -1,73 +1,20 @@
 import json
 import time
 from collections import defaultdict
-from typing import Any, BinaryIO, Dict, Iterable, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
 import onnxruntime as ort
-import soundfile as sf
 import torch
 from lhotse import FbankConfig
-from lhotse.augmentation import get_or_create_resampler
 from lhotse.features.kaldi.layers import Wav2LogFilterBank
 from lhotse.utils import Pathlike
 
-from lattifai.errors import AlignmentError, AudioFormatError, AudioLoadError, DependencyError, ModelLoadError
+from lattifai.audio2 import AudioData
+from lattifai.errors import AlignmentError, DependencyError, ModelLoadError
 
-ChannelSelectorType = Union[int, Iterable[int], str]
 
-
-def resample_audio(
-    audio_sr: Tuple[torch.Tensor, int],
-    sampling_rate: int,
-    device: Optional[str],
-    channel_selector: Optional[ChannelSelectorType] = "average",
-) -> torch.Tensor:
-    """
-    return:
-        (1, T)
-    """
-    audio, sr = audio_sr
-
-    if channel_selector is None:
-        # keep the original multi-channel signal
-        tensor = audio
-    elif isinstance(channel_selector, int):
-        assert audio.shape[0] >= channel_selector, f"Invalid channel: {channel_selector}"
-        tensor = audio[channel_selector : channel_selector + 1].clone()
-        del audio
-    elif isinstance(channel_selector, str):
-        assert channel_selector == "average"
-        tensor = torch.mean(audio.to(device), dim=0, keepdim=True)
-        del audio
-    else:
-        assert isinstance(channel_selector, Iterable)
-        num_channels = audio.shape[0]
-        print(f"Selecting channels {channel_selector} from the signal with {num_channels} channels.")
-        assert isinstance(channel_selector, Iterable)
-        if max(channel_selector) >= num_channels:
-            raise ValueError(
-                f"Cannot select channel subset {channel_selector} from a signal with {num_channels} channels."
-            )
-        tensor = audio[channel_selector]
-
-    tensor = tensor.to(device)
-    if sr != sampling_rate:
-        resampler = get_or_create_resampler(sr, sampling_rate).to(device=device)
-        length = tensor.size(-1)
-        chunk_size = sampling_rate * 3600
-        if length > chunk_size:
-            resampled_chunks = []
-            for i in range(0, length, chunk_size):
-                resampled_chunks.append(resampler(tensor[..., i : i + chunk_size]))
-            tensor = torch.cat(resampled_chunks, dim=-1)
-        else:
-            tensor = resampler(tensor)
-
-    return tensor
-
-
-class Lattice1AlphaWorker:
+class Lattice1Worker:
     """Worker for processing audio with LatticeGraph."""
 
     def __init__(self, model_path: Pathlike, device: str = "cpu", num_threads: int = 8) -> None:
@@ -109,6 +56,10 @@ class Lattice1AlphaWorker:
         self.device = torch.device(device)
         self.timings = defaultdict(lambda: 0.0)
 
+    @property
+    def frame_shift(self) -> float:
+        return 0.02  # 20 ms
+
     @torch.inference_mode()
     def emission(self, audio: torch.Tensor) -> torch.Tensor:
         _start = time.time()
@@ -138,68 +89,17 @@ class Lattice1AlphaWorker:
         self.timings["emission"] += time.time() - _start
         return emission  # (1, T, vocab_size) torch
 
-    def load_audio(
-        self, audio: Union[Pathlike, BinaryIO], channel_selector: Optional[ChannelSelectorType] = "average"
-    ) -> Tuple[torch.Tensor, int]:
-        # load audio
-        try:
-            waveform, sample_rate = sf.read(audio, always_2d=True, dtype="float32")  # numpy array
-            waveform = waveform.T  # (channels, samples)
-        except Exception as primary_error:
-            # Fallback to PyAV for formats not supported by soundfile
-            try:
-                import av
-            except ImportError:
-                raise DependencyError(
-                    "av (PyAV)", install_command="pip install av", context={"primary_error": str(primary_error)}
-                )
-
-            try:
-                container = av.open(audio)
-                audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-
-                if audio_stream is None:
-                    raise AudioFormatError(str(audio), "No audio stream found in file")
-
-                # Resample to target sample rate during decoding
-                audio_stream.codec_context.format = av.AudioFormat("flt")  # 32-bit float
-
-                frames = []
-                for frame in container.decode(audio_stream):
-                    # Convert frame to numpy array
-                    array = frame.to_ndarray()
-                    # Ensure shape is (channels, samples)
-                    if array.ndim == 1:
-                        array = array.reshape(1, -1)
-                    elif array.ndim == 2 and array.shape[0] > array.shape[1]:
-                        array = array.T
-                    frames.append(array)
-
-                container.close()
-
-                if not frames:
-                    raise AudioFormatError(str(audio), "No audio data found in file")
-
-                # Concatenate all frames
-                waveform = np.concatenate(frames, axis=1)
-                sample_rate = audio_stream.codec_context.sample_rate
-            except Exception as e:
-                raise AudioLoadError(str(audio), original_error=e)
-
-        return resample_audio(
-            (torch.from_numpy(waveform), sample_rate),
-            self.config.get("sampling_rate", 16000),
-            device=self.device.type,
-            channel_selector=channel_selector,
-        )
-
     def alignment(
-        self, audio: Union[Union[Pathlike, BinaryIO], torch.tensor], lattice_graph: Tuple[str, int, float]
+        self,
+        audio: AudioData,
+        lattice_graph: Tuple[str, int, float],
+        emission: Optional[torch.Tensor] = None,
+        offset: float = 0.0,
     ) -> Dict[str, Any]:
         """Process audio with LatticeGraph.
 
         Args:
-            audio: Audio file path or binary data
+            audio: AudioData object
             lattice_graph: LatticeGraph data
 
         Returns:
@@ -210,22 +110,15 @@ class Lattice1AlphaWorker:
             DependencyError: If required dependencies are missing
             AlignmentError: If alignment process fails
         """
-        # load audio
-        if isinstance(audio, torch.Tensor):
-            waveform = audio
-        else:
-            waveform = self.load_audio(audio)  # (1, L)
-
-        _start = time.time()
-        try:
-            emission = self.emission(waveform.to(self.device))  # (1, T, vocab_size)
-        except Exception as e:
-            raise AlignmentError(
-                "Failed to compute acoustic features from audio",
-                audio_path=str(audio) if not isinstance(audio, torch.Tensor) else "tensor",
-                context={"original_error": str(e)},
-            )
-        self.timings["emission"] += time.time() - _start
+        if emission is None:
+            try:
+                emission = self.emission(audio.tensor.to(self.device))  # (1, T, vocab_size)
+            except Exception as e:
+                raise AlignmentError(
+                    "Failed to compute acoustic features from audio",
+                    media_path=str(audio) if not isinstance(audio, torch.Tensor) else "tensor",
+                    context={"original_error": str(e)},
+                )
 
         try:
             import k2
@@ -275,10 +168,18 @@ class Lattice1AlphaWorker:
         except Exception as e:
             raise AlignmentError(
                 "Failed to perform forced alignment",
-                audio_path=str(audio) if not isinstance(audio, torch.Tensor) else "tensor",
+                media_path=str(audio) if not isinstance(audio, torch.Tensor) else "tensor",
                 context={"original_error": str(e), "emission_shape": list(emission.shape), "device": str(device)},
             )
         self.timings["align_segments"] += time.time() - _start
 
         channel = 0
-        return emission, results, labels, 0.02, 0.0, channel  # frame_shift=20ms, offset=0.0s
+        return emission, results, labels, self.frame_shift, offset, channel  # frame_shift=20ms
+
+
+def _load_worker(model_path: str, device: str) -> Lattice1Worker:
+    """Instantiate lattice worker with consistent error handling."""
+    try:
+        return Lattice1Worker(model_path, device=device, num_threads=8)
+    except Exception as e:
+        raise ModelLoadError(f"worker from {model_path}", original_error=e)

lattifai/{tokenizer → alignment}/phonemizer.py

@@ -1,7 +1,7 @@
 import re
 from typing import List, Optional, Union
 
-from dp.phonemizer import Phonemizer  # g2p-phonemizer
+from g2pp.phonemizer import Phonemizer  # g2p-phonemizer
 from num2words import num2words
 
 LANGUAGE = "omni"

lattifai/alignment/segmenter.py

@@ -0,0 +1,166 @@
+"""Segmented alignment for long audio files."""
+
+from typing import List, Optional, Tuple
+
+import colorful
+
+from lattifai.audio2 import AudioData
+from lattifai.caption import Caption, Supervision
+from lattifai.config import AlignmentConfig
+
+from .tokenizer import END_PUNCTUATION
+
+
+class Segmenter:
+    """
+    Handles segmented alignment for long audio/video files.
+
+    Instead of aligning the entire audio at once (which can be slow and memory-intensive
+    for long files), this class splits the alignment into manageable segments based on
+    caption boundaries, time intervals, or an adaptive strategy.
+    """
+
+    def __init__(self, config: AlignmentConfig):
+        """
+        Initialize segmented aligner.
+
+        Args:
+            config: Alignment configuration with segmentation parameters
+        """
+        self.config = config
+
+    def __call__(
+        self,
+        caption: Caption,
+        max_duration: Optional[float] = None,
+    ) -> List[Tuple[float, float, List[Supervision]]]:
+        """
+        Create segments based on caption boundaries and gaps.
+
+        Splits when:
+        1. Gap between captions exceeds segment_max_gap
+        2. Duration approaches max_duration (adaptive mode only) and there's a reasonable break
+        3. Duration significantly exceeds max_duration (adaptive mode only)
+
+        Args:
+            caption: Caption object with supervisions
+            max_duration: Optional maximum segment duration (enables adaptive behavior)
+
+        Returns:
+            List of (start_time, end_time, supervisions) tuples for each segment
+        """
+        if not max_duration:
+            max_duration = self.config.segment_duration
+
+        if not caption.supervisions:
+            return []
+
+        supervisions = sorted(caption.supervisions, key=lambda s: s.start)
+
+        segments = []
+        current_segment_sups = []
+
+        def should_skipalign(sups):
+            return len(sups) == 1 and sups[0].text.strip().startswith("[") and sups[0].text.strip().endswith("]")
+
+        for i, sup in enumerate(supervisions):
+            if not current_segment_sups:
+                current_segment_sups.append(sup)
+                if should_skipalign(current_segment_sups):
+                    # Single [APPLAUSE] caption, make its own segment
+                    segments.append(
+                        (current_segment_sups[0].start, current_segment_sups[-1].end, current_segment_sups, True)
+                    )
+                    current_segment_sups = []
+                continue
+
+            prev_sup = supervisions[i - 1]
+
+            gap = max(sup.start - prev_sup.end, 0.0)
+            # Always split on large gaps (natural breaks)
+            exclude_max_gap = False
+            if gap > self.config.segment_max_gap:
+                exclude_max_gap = True
+
+            endswith_punc = any(sup.text.endswith(punc) for punc in END_PUNCTUATION)
+
+            # Adaptive duration control
+            segment_duration = sup.end - current_segment_sups[0].start
+
+            # Split if approaching duration limit and there's a reasonable break
+            should_split = False
+            if segment_duration >= max_duration * 0.8 and gap >= 1.0:
+                should_split = True
+
+            # Force split if duration exceeded significantly
+            exclude_max_duration = False
+            if segment_duration >= max_duration * 1.2:
+                exclude_max_duration = True
+
+            # [APPLAUSE] [APPLAUSE] [MUSIC]
+            if sup.text.strip().startswith("[") and sup.text.strip().endswith("]"):
+                # Close current segment
+                if current_segment_sups:
+                    segment_start = current_segment_sups[0].start
+                    segment_end = current_segment_sups[-1].end + min(gap / 2.0, 2.0)
+                    segments.append(
+                        (segment_start, segment_end, current_segment_sups, should_skipalign(current_segment_sups))
+                    )
+
+                # Add current supervision as its own segment
+                segments.append((sup.start, sup.end, [sup], True))
+
+                # Update reset for new segment
+                current_segment_sups = []
+                continue
+
+            if (should_split and endswith_punc) or exclude_max_gap or exclude_max_duration:
+                # Close current segment
+                if current_segment_sups:
+                    segment_start = current_segment_sups[0].start
+                    segment_end = current_segment_sups[-1].end + min(gap / 2.0, 2.0)
+                    segments.append(
+                        (segment_start, segment_end, current_segment_sups, should_skipalign(current_segment_sups))
+                    )
+
+                # Start new segment
+                current_segment_sups = [sup]
+            else:
+                current_segment_sups.append(sup)
+
+        # Add final segment
+        if current_segment_sups:
+            segment_start = current_segment_sups[0].start
+            segment_end = current_segment_sups[-1].end + 2.0
+            segments.append((segment_start, segment_end, current_segment_sups, should_skipalign(current_segment_sups)))
+
+        return segments
+
+    def print_segment_info(
+        self,
+        segments: List[Tuple[float, float, List[Supervision]]],
+        verbose: bool = True,
+    ) -> None:
+        """
+        Print information about created segments.
+
+        Args:
+            segments: List of segment tuples
+            verbose: Whether to print detailed info
+        """
+        if not verbose:
+            return
+
+        total_sups = sum(len(sups) if isinstance(sups, list) else 1 for _, _, sups, _ in segments)
+
+        print(colorful.cyan(f"📊 Created {len(segments)} alignment segments:"))
+        for i, (start, end, sups, _) in enumerate(segments, 1):
+            duration = end - start
+            print(
+                colorful.white(
+                    f"   Segment {i:04d}: {start:8.2f}s - {end:8.2f}s "
+                    f"(duration: {duration:8.2f}s, supervisions: {len(sups)if isinstance(sups, list) else 1:4d})"
+                )
+            )
+
+        print(colorful.green(f"   Total: {total_sups} supervisions across {len(segments)} segments"))