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

lattifai/__init__.py

@@ -1,34 +1,45 @@
-import os
 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')
+    __version__ = version("lattifai")
 except Exception:
-    __version__ = '0.1.0'  # fallback version
+    __version__ = "0.1.0"  # fallback version
 
 
 # Check and auto-install k2 if not present
@@ -39,15 +50,15 @@ def _check_and_install_k2():
     except ImportError:
         import subprocess
 
-        print('k2 is not installed. Attempting to install k2...')
+        print("k2 is not installed. Attempting to install k2...")
         try:
-            subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'install-k2'])
-            subprocess.check_call([sys.executable, '-m', 'install_k2'])
+            subprocess.check_call([sys.executable, "-m", "pip", "install", "install-k2"])
+            subprocess.check_call([sys.executable, "-m", "install_k2"])
             import k2  # Try importing again after installation
 
-            print('k2 installed successfully.')
+            print("k2 installed successfully.")
         except Exception as e:
-            warnings.warn(f'Failed to install k2 automatically. Please install it manually. Error: {e}')
+            warnings.warn(f"Failed to install k2 automatically. Please install it manually. Error: {e}")
     return True
 
 
@@ -55,35 +66,38 @@ 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
-    'LattifAIError',
-    'AudioProcessingError',
-    'AudioLoadError',
-    'AudioFormatError',
-    'SubtitleProcessingError',
-    'SubtitleParseError',
-    'AlignmentError',
-    'LatticeEncodingError',
-    'LatticeDecodingError',
-    'ModelLoadError',
-    'DependencyError',
-    'APIError',
-    'ConfigurationError',
-    'SubtitleIO',
-    '__version__',
+    # Client classes
+    "LattifAI",
+    # Config classes
+    "AlignmentConfig",
+    "ClientConfig",
+    "CaptionConfig",
+    "DiarizationConfig",
+    "MediaConfig",
+    "AUDIO_FORMATS",
+    "VIDEO_FORMATS",
+    "MEDIA_FORMATS",
+    # Error classes
+    "LattifAIError",
+    "AudioProcessingError",
+    "AudioLoadError",
+    "AudioFormatError",
+    "CaptionProcessingError",
+    "CaptionParseError",
+    "AlignmentError",
+    "LatticeEncodingError",
+    "LatticeDecodingError",
+    "ModelLoadError",
+    "DependencyError",
+    "APIError",
+    "ConfigurationError",
+    # 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/alignment/lattice1_worker.py

@@ -0,0 +1,185 @@
+import json
+import time
+from collections import defaultdict
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import onnxruntime as ort
+import torch
+from lhotse import FbankConfig
+from lhotse.features.kaldi.layers import Wav2LogFilterBank
+from lhotse.utils import Pathlike
+
+from lattifai.audio2 import AudioData
+from lattifai.errors import AlignmentError, DependencyError, ModelLoadError
+
+
+class Lattice1Worker:
+    """Worker for processing audio with LatticeGraph."""
+
+    def __init__(self, model_path: Pathlike, device: str = "cpu", num_threads: int = 8) -> None:
+        try:
+            self.config = json.load(open(f"{model_path}/config.json"))
+        except Exception as e:
+            raise ModelLoadError(f"config from {model_path}", original_error=e)
+
+        # SessionOptions
+        sess_options = ort.SessionOptions()
+        # sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.intra_op_num_threads = num_threads  # CPU cores
+        sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
+        sess_options.add_session_config_entry("session.intra_op.allow_spinning", "0")
+
+        providers = []
+        if device.startswith("cuda") and ort.get_all_providers().count("CUDAExecutionProvider") > 0:
+            providers.append("CUDAExecutionProvider")
+        elif device.startswith("mps") and ort.get_all_providers().count("MPSExecutionProvider") > 0:
+            providers.append("MPSExecutionProvider")
+
+        try:
+            self.acoustic_ort = ort.InferenceSession(
+                f"{model_path}/acoustic_opt.onnx",
+                sess_options,
+                providers=providers + ["CPUExecutionProvider", "CoreMLExecutionProvider"],
+            )
+        except Exception as e:
+            raise ModelLoadError(f"acoustic model from {model_path}", original_error=e)
+
+        try:
+            config = FbankConfig(num_mel_bins=80, device=device, snip_edges=False)
+            config_dict = config.to_dict()
+            config_dict.pop("device")
+            self.extractor = Wav2LogFilterBank(**config_dict).to(device).eval()
+        except Exception as e:
+            raise ModelLoadError(f"feature extractor for device {device}", original_error=e)
+
+        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()
+        # audio -> features -> emission
+        features = self.extractor(audio)  # (1, T, D)
+        if features.shape[1] > 6000:
+            features_list = torch.split(features, 6000, dim=1)
+            emissions = []
+            for features in features_list:
+                ort_inputs = {
+                    "features": features.cpu().numpy(),
+                    "feature_lengths": np.array([features.size(1)], dtype=np.int64),
+                }
+                emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+                emissions.append(emission)
+            emission = torch.cat(
+                [torch.from_numpy(emission).to(self.device) for emission in emissions], dim=1
+            )  # (1, T, vocab_size)
+        else:
+            ort_inputs = {
+                "features": features.cpu().numpy(),
+                "feature_lengths": np.array([features.size(1)], dtype=np.int64),
+            }
+            emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+            emission = torch.from_numpy(emission).to(self.device)
+
+        self.timings["emission"] += time.time() - _start
+        return emission  # (1, T, vocab_size) torch
+
+    def alignment(
+        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: AudioData object
+            lattice_graph: LatticeGraph data
+
+        Returns:
+            Processed LatticeGraph
+
+        Raises:
+            AudioLoadError: If audio cannot be loaded
+            DependencyError: If required dependencies are missing
+            AlignmentError: If alignment process fails
+        """
+        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
+        except ImportError:
+            raise DependencyError("k2", install_command="pip install install-k2 && python -m install_k2")
+
+        try:
+            from lattifai_core.lattice.decode import align_segments
+        except ImportError:
+            raise DependencyError("lattifai_core", install_command="Contact support for lattifai_core installation")
+
+        lattice_graph_str, final_state, acoustic_scale = lattice_graph
+
+        _start = time.time()
+        try:
+            # graph
+            decoding_graph = k2.Fsa.from_str(lattice_graph_str, acceptor=False)
+            decoding_graph.requires_grad_(False)
+            decoding_graph = k2.arc_sort(decoding_graph)
+            decoding_graph.skip_id = int(final_state)
+            decoding_graph.return_id = int(final_state + 1)
+        except Exception as e:
+            raise AlignmentError(
+                "Failed to create decoding graph from lattice",
+                context={"original_error": str(e), "lattice_graph_length": len(lattice_graph_str)},
+            )
+        self.timings["decoding_graph"] += time.time() - _start
+
+        _start = time.time()
+        if self.device.type == "mps":
+            device = "cpu"  # k2 does not support mps yet
+        else:
+            device = self.device
+
+        try:
+            results, labels = align_segments(
+                emission.to(device) * acoustic_scale,
+                decoding_graph.to(device),
+                torch.tensor([emission.shape[1]], dtype=torch.int32),
+                search_beam=200,
+                output_beam=80,
+                min_active_states=400,
+                max_active_states=10000,
+                subsampling_factor=1,
+                reject_low_confidence=False,
+            )
+        except Exception as e:
+            raise AlignmentError(
+                "Failed to perform forced alignment",
+                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, 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,16 +1,16 @@
 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'
+LANGUAGE = "omni"
 
 
 class G2Phonemizer:
     def __init__(self, model_checkpoint, device):
         self.phonemizer = Phonemizer.from_checkpoint(model_checkpoint, device=device).predictor
-        self.pattern = re.compile(r'\d+')
+        self.pattern = re.compile(r"\d+")
 
     def num2words(self, word, lang: str):
         matches = self.pattern.findall(word)
@@ -31,7 +31,7 @@ class G2Phonemizer:
             is_list = False
 
         predictions = self.phonemizer(
-            [self.num2words(word.replace(' .', '.').replace('.', ' .'), lang=lang or 'en') for word in words],
+            [self.num2words(word.replace(" .", ".").replace(".", " ."), lang=lang or "en") for word in words],
             lang=LANGUAGE,
             batch_size=min(batch_size or len(words), 128),
             num_prons=num_prons,

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"))