lattifai 1.2.2__py3-none-any.whl → 1.3.1__py3-none-any.whl

lattifai/data/caption.py

@@ -0,0 +1,228 @@
+"""Extended Caption class with transcription, alignment, and diarization support."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar
+
+from lattifai.caption import Caption as BaseCaption
+from lattifai.caption import Pathlike, Supervision
+
+if TYPE_CHECKING:
+    from lattifai_core.event import LEDOutput
+
+DiarizationOutput = TypeVar("DiarizationOutput")
+
+
+@dataclass
+class Caption(BaseCaption):
+    """
+    Extended Caption with transcription, alignment, and diarization support.
+
+    Inherits from BaseCaption and adds fields for:
+    - alignments: Post-alignment results
+    - transcription: ASR results
+    - event: LattifAI Event Detection results (LEDOutput)
+    - diarization: Speaker diarization results
+
+
+    These fields are used in the LattifAI pipeline for:
+    - Forced alignment results
+    - Storing intermediate transcription results
+    - LattifAI Event Detection (music, applause, speech, etc.)
+    - Speaker identification and separation
+
+    """
+
+    # Alignment results
+    alignments: List[Supervision] = field(default_factory=list)
+
+    # Transcription results
+    transcription: List[Supervision] = field(default_factory=list)
+
+    # LattifAI Event Detection results
+    event: Optional["LEDOutput"] = None
+
+    # Speaker Diarization results
+    diarization: Optional[DiarizationOutput] = None
+
+    def __len__(self) -> int:
+        """Return the number of supervision segments."""
+        return len(self.supervisions or self.transcription)
+
+    def __repr__(self) -> str:
+        """String representation of Caption."""
+        lang = f"lang={self.language}" if self.language else "lang=unknown"
+        kind_str = f"kind={self.kind}" if self.kind else ""
+        parts = [f"Caption({len(self.supervisions or self.transcription)} segments", lang]
+        if kind_str:
+            parts.append(kind_str)
+        if self.duration:
+            parts.append(f"duration={self.duration:.2f}s")
+        return ", ".join(parts) + ")"
+
+    def with_margins(
+        self,
+        start_margin: float = 0.08,
+        end_margin: float = 0.20,
+        min_gap: float = 0.08,
+        collision_mode: str = "trim",
+    ) -> "Caption":
+        """
+        Create a new Caption with segment boundaries adjusted based on word-level alignment.
+
+        Uses supervision.alignment['word'] to recalculate segment start/end times
+        with the specified margins applied around the actual speech boundaries.
+
+        Prefers alignments > supervisions > transcription as source.
+
+        Args:
+            start_margin: Seconds to extend before the first word (default: 0.08)
+            end_margin: Seconds to extend after the last word (default: 0.20)
+            min_gap: Minimum gap between segments for collision handling (default: 0.08)
+            collision_mode: How to handle segment overlap - 'trim' or 'gap' (default: 'trim')
+
+        Returns:
+            New Caption instance with adjusted timestamps
+
+        Note:
+            Segments without alignment data will keep their original timestamps.
+        """
+        from lattifai.caption.standardize import apply_margins_to_captions
+
+        # Determine which supervisions to use (priority: alignments > supervisions > transcription)
+        if self.alignments:
+            source_sups = self.alignments
+        elif self.supervisions:
+            source_sups = self.supervisions
+        else:
+            source_sups = self.transcription
+
+        adjusted_sups = apply_margins_to_captions(
+            source_sups,
+            start_margin=start_margin,
+            end_margin=end_margin,
+            min_gap=min_gap,
+            collision_mode=collision_mode,
+        )
+
+        return Caption(
+            supervisions=adjusted_sups,
+            transcription=self.transcription,
+            event=self.event,
+            diarization=self.diarization,
+            alignments=[],  # Clear alignments since we've applied them
+            language=self.language,
+            kind=self.kind,
+            source_format=self.source_format,
+            source_path=self.source_path,
+            metadata=self.metadata.copy() if self.metadata else {},
+        )
+
+    def write(
+        self,
+        path=None,
+        output_format: Optional[str] = None,
+        include_speaker_in_text: bool = True,
+        word_level: bool = False,
+        karaoke_config=None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Write caption to file or return as bytes.
+
+        Prefers alignments > supervisions > transcription as source.
+
+        Args:
+            path: Path to output caption file, BytesIO object, or None to return bytes
+            output_format: Output format (e.g., 'srt', 'vtt', 'ass')
+            include_speaker_in_text: Whether to include speaker labels in text
+            word_level: Use word-level output format if supported
+            karaoke_config: Karaoke configuration
+            metadata: Optional metadata dict to pass to writer
+
+        Returns:
+            Path to the written file if path is a file path, or bytes if path is BytesIO/None
+        """
+        # Temporarily swap supervisions to use the priority order
+        original_supervisions = self.supervisions
+
+        if self.alignments:
+            self.supervisions = self.alignments
+        elif not self.supervisions and self.transcription:
+            self.supervisions = self.transcription
+
+        try:
+            result = super().write(
+                path=path,
+                output_format=output_format,
+                include_speaker_in_text=include_speaker_in_text,
+                word_level=word_level,
+                karaoke_config=karaoke_config,
+                metadata=metadata,
+            )
+        finally:
+            # Restore original supervisions
+            self.supervisions = original_supervisions
+
+        return result
+
+    @classmethod
+    def from_transcription_results(
+        cls,
+        transcription: List[Supervision],
+        event: Optional["LEDOutput"] = None,
+        diarization: Optional[DiarizationOutput] = None,
+        language: Optional[str] = None,
+        source_path: Optional[Pathlike] = None,
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> "Caption":
+        """
+        Create Caption from transcription results including audio events and diarization.
+
+        Args:
+            transcription: List of transcription supervision segments
+            event: Optional LEDOutput with event detection results
+            diarization: Optional DiarizationOutput with speaker diarization results
+            language: Language code
+            source_path: Source file path
+            metadata: Additional metadata
+
+        Returns:
+            New Caption instance with transcription data
+        """
+        return cls(
+            transcription=transcription,
+            event=event,
+            diarization=diarization,
+            language=language,
+            kind="transcription",
+            source_format="asr",
+            source_path=source_path,
+            metadata=metadata or {},
+        )
+
+    def read_diarization(
+        self,
+        path: Pathlike,
+    ) -> "DiarizationOutput":
+        """
+        Read speaker diarization TextGrid from file.
+        """
+        from lattifai_core.diarization import DiarizationOutput
+
+        self.diarization = DiarizationOutput.read(path)
+        return self.diarization
+
+    def write_diarization(
+        self,
+        path: Pathlike,
+    ) -> Pathlike:
+        """
+        Write speaker diarization TextGrid to file.
+        """
+        if not self.diarization:
+            raise ValueError("No speaker diarization data to write.")
+
+        self.diarization.write(path)
+        return path

lattifai/errors.py

@@ -1,10 +1,42 @@
 """Error handling and exception classes for LattifAI SDK."""
 
+import functools
 import traceback
 from typing import Any, Dict, Optional
 
 import colorful
 
+
+def format_exception(e: "LattifAIError") -> str:
+    """Format LattifAIError with filtered traceback (only lattifai frames)."""
+    tb_lines = traceback.format_exception(type(e), e, e.__traceback__)
+    filtered = []
+    skip_next_code_line = False
+
+    for i, line in enumerate(tb_lines):
+        if skip_next_code_line:
+            skip_next_code_line = False
+            continue
+
+        if line.startswith("Traceback") or not line.startswith("  File"):
+            filtered.append(line)
+        elif "lattifai" in line:
+            filtered.append(line)
+            if i + 1 < len(tb_lines) and tb_lines[i + 1].startswith("    "):
+                filtered.append(tb_lines[i + 1])
+                skip_next_code_line = True
+        elif i + 1 < len(tb_lines) and tb_lines[i + 1].startswith("    "):
+            skip_next_code_line = True
+
+    return "".join(filtered)
+
+
+def _merge_context(kwargs: Dict[str, Any], updates: Dict[str, Any]) -> None:
+    """Merge updates into kwargs['context'], creating it if needed."""
+    context = kwargs.setdefault("context", {})
+    context.update(updates)
+
+
 # Error help messages
 LATTICE_DECODING_FAILURE_HELP = (
     "Failed to decode lattice alignment. Possible reasons:\n\n"
@@ -76,10 +108,8 @@ class AudioProcessingError(LattifAIError):
     """Error during audio processing operations."""
 
     def __init__(self, message: str, media_path: Optional[str] = None, **kwargs):
-        context = kwargs.get("context", {})
         if media_path:
-            context["media_path"] = media_path
-        kwargs["context"] = context
+            _merge_context(kwargs, {"media_path": media_path})
         super().__init__(message, **kwargs)
 
 
@@ -90,11 +120,9 @@ class AudioLoadError(AudioProcessingError):
         message = f"Failed to load audio file: {colorful.red(media_path)}"
         if original_error:
             message += f" - {colorful.red(str(original_error))}"
-
-        context = kwargs.get("context", {})
-        context.update({"media_path": media_path, "original_error": str(original_error) if original_error else None})
-        kwargs["context"] = context
-
+        _merge_context(
+            kwargs, {"media_path": media_path, "original_error": str(original_error) if original_error else None}
+        )
         super().__init__(message, media_path=media_path, **kwargs)
 
 
@@ -103,9 +131,7 @@ class AudioFormatError(AudioProcessingError):
 
     def __init__(self, media_path: str, format_issue: str, **kwargs):
         message = f"Audio format error for {colorful.red(media_path)}: {colorful.red(format_issue)}"
-        context = kwargs.get("context", {})
-        context.update({"media_path": media_path, "format_issue": format_issue})
-        kwargs["context"] = context
+        _merge_context(kwargs, {"media_path": media_path, "format_issue": format_issue})
         super().__init__(message, media_path=media_path, **kwargs)
 
 
@@ -113,10 +139,8 @@ class CaptionProcessingError(LattifAIError):
     """Error during caption/text processing operations."""
 
     def __init__(self, message: str, caption_path: Optional[str] = None, **kwargs):
-        context = kwargs.get("context", {})
         if caption_path:
-            context["caption_path"] = caption_path
-        kwargs["context"] = context
+            _merge_context(kwargs, {"caption_path": caption_path})
         super().__init__(message, **kwargs)
 
 
@@ -125,9 +149,7 @@ class CaptionParseError(CaptionProcessingError):
 
     def __init__(self, caption_path: str, parse_issue: str, **kwargs):
         message = f"Failed to parse caption file {caption_path}: {parse_issue}"
-        context = kwargs.get("context", {})
-        context.update({"caption_path": caption_path, "parse_issue": parse_issue})
-        kwargs["context"] = context
+        _merge_context(kwargs, {"caption_path": caption_path, "parse_issue": parse_issue})
         super().__init__(message, caption_path=caption_path, **kwargs)
 
 
@@ -135,12 +157,13 @@ class AlignmentError(LattifAIError):
     """Error during audio-text alignment process."""
 
     def __init__(self, message: str, media_path: Optional[str] = None, caption_path: Optional[str] = None, **kwargs):
-        context = kwargs.get("context", {})
+        updates = {}
         if media_path:
-            context["media_path"] = media_path
+            updates["media_path"] = media_path
         if caption_path:
-            context["caption_path"] = caption_path
-        kwargs["context"] = context
+            updates["caption_path"] = caption_path
+        if updates:
+            _merge_context(kwargs, updates)
         super().__init__(message, **kwargs)
 
 
@@ -151,36 +174,44 @@ class LatticeEncodingError(AlignmentError):
         message = "Failed to generate lattice graph from text"
         if original_error:
             message += f": {colorful.red(str(original_error))}"
-
-        context = kwargs.get("context", {})
-        context.update(
+        text_preview = text_content[:100] + "..." if len(text_content) > 100 else text_content
+        _merge_context(
+            kwargs,
             {
                 "text_content_length": len(text_content),
-                "text_preview": text_content[:100] + "..." if len(text_content) > 100 else text_content,
+                "text_preview": text_preview,
                 "original_error": str(original_error) if original_error else None,
-            }
+            },
         )
-        kwargs["context"] = context
         super().__init__(message, **kwargs)
 
 
 class LatticeDecodingError(AlignmentError):
     """Error decoding lattice alignment results."""
 
-    def __init__(self, lattice_id: str, original_error: Optional[Exception] = None, **kwargs):
-        message = f"Failed to decode lattice alignment results for lattice ID: {colorful.red(lattice_id)}"
+    def __init__(
+        self,
+        lattice_id: str,
+        message: Optional[str] = None,
+        original_error: Optional[Exception] = None,
+        skip_help: bool = False,
+        **kwargs,
+    ):
+        message = message or f"Failed to decode lattice alignment results for lattice ID: {colorful.red(lattice_id)}"
 
-        # Don't duplicate the help message if it's already in original_error
-        if original_error and str(original_error) != LATTICE_DECODING_FAILURE_HELP:
-            message += f" - {colorful.red(str(original_error))}"
+        error_str = str(original_error) if original_error else None
+        is_help_message = error_str == LATTICE_DECODING_FAILURE_HELP
+
+        if original_error and not is_help_message:
+            message += f" - {colorful.red(error_str)}"
+
+        context_updates = {"lattice_id": lattice_id}
+        if original_error and not is_help_message:
+            context_updates["original_error"] = error_str
+        _merge_context(kwargs, context_updates)
 
-        context = kwargs.get("context", {})
-        # Don't store the entire help message in context to avoid duplication
-        if original_error and str(original_error) != LATTICE_DECODING_FAILURE_HELP:
-            context["original_error"] = str(original_error)
-        context["lattice_id"] = lattice_id
-        kwargs["context"] = context
         super().__init__(message, **kwargs)
+        self.skip_help = skip_help
 
     def get_message(self) -> str:
         """Return formatted error message with help text."""
@@ -188,8 +219,9 @@ class LatticeDecodingError(AlignmentError):
         if self.context and self.context.get("lattice_id"):
             # Only show essential context (lattice_id), not the duplicated help message
             base_message += f'\n{colorful.yellow("Lattice ID:")} {self.context["lattice_id"]}'
-        # Append help message once at the end
-        base_message += f"\n\n{colorful.yellow(LATTICE_DECODING_FAILURE_HELP)}"
+        # Append help message only if not skipped (e.g., when anomaly info is provided)
+        if not self.skip_help:
+            base_message += f"\n\n{colorful.yellow(LATTICE_DECODING_FAILURE_HELP)}"
         return base_message
 
 
@@ -200,10 +232,9 @@ class ModelLoadError(LattifAIError):
         message = f"Failed to load model: {colorful.red(model_name)}"
         if original_error:
             message += f" - {colorful.red(str(original_error))}"
-
-        context = kwargs.get("context", {})
-        context.update({"model_name": model_name, "original_error": str(original_error) if original_error else None})
-        kwargs["context"] = context
+        _merge_context(
+            kwargs, {"model_name": model_name, "original_error": str(original_error) if original_error else None}
+        )
         super().__init__(message, **kwargs)
 
 
@@ -214,10 +245,7 @@ class DependencyError(LattifAIError):
         message = f"Missing required dependency: {colorful.red(dependency_name)}"
         if install_command:
             message += f"\nPlease install it using: {colorful.yellow(install_command)}"
-
-        context = kwargs.get("context", {})
-        context.update({"dependency_name": dependency_name, "install_command": install_command})
-        kwargs["context"] = context
+        _merge_context(kwargs, {"dependency_name": dependency_name, "install_command": install_command})
         super().__init__(message, **kwargs)
 
 
@@ -225,9 +253,7 @@ class APIError(LattifAIError):
     """Error communicating with LattifAI API."""
 
     def __init__(self, message: str, status_code: Optional[int] = None, response_text: Optional[str] = None, **kwargs):
-        context = kwargs.get("context", {})
-        context.update({"status_code": status_code, "response_text": response_text})
-        kwargs["context"] = context
+        _merge_context(kwargs, {"status_code": status_code, "response_text": response_text})
         super().__init__(message, **kwargs)
 
 
@@ -249,14 +275,13 @@ class QuotaExceededError(APIError):
 def handle_exception(func):
     """Decorator to handle exceptions and convert them to LattifAI errors."""
 
+    @functools.wraps(func)
     def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except LattifAIError:
-            # Re-raise LattifAI errors as-is
             raise
         except Exception as e:
-            # Convert other exceptions to LattifAI errors
             error_msg = f"Unexpected error in {func.__name__}: {str(e)}"
             context = {
                 "function": func.__name__,

lattifai/event/__init__.py

@@ -0,0 +1,65 @@
+"""Audio Event Detection module for LattifAI.
+
+This module provides audio event detection capabilities, it can identify various
+audio events including speech, music, singing, and demographic characteristics
+(male, female, child voices).
+
+Key Components:
+    LattifAIEventDetector: Main class that wraps lattifai_core's
+        EventDetector for seamless integration with LattifAI workflows.
+
+Features:
+    - Multi-class audio event detection (30+ reduced classes or 400+ full classes)
+    - Voice Activity Detection (VAD) for speech segmentation
+    - Gender/age classification for speech segments
+    - Configurable detection thresholds and top-k filtering
+    - Support for both bundled and custom pretrained models
+
+Detected Event Types:
+    - Speech: General speech activity
+    - Male/Female/Child: Speaker demographic classification
+    - Music: Musical content detection
+    - Singing: Vocal music detection
+    - Synthetic: Synthetic/electronic sounds
+
+Configuration:
+    Use EventConfig to control:
+    - enabled: Whether to run audio event detection
+    - device: GPU/CPU device selection
+    - dtype: Model precision (float32, float16, bfloat16)
+    - reduced: Use reduced label set (33 vs 400+ classes)
+    - top_k: Number of top event classes to detect
+    - vad_chunk_size/vad_max_gap: VAD segmentation parameters
+
+Example:
+    >>> from lattifai.event import LattifAIEventDetector
+    >>> from lattifai.config import EventConfig
+    >>> from lattifai.audio2 import AudioLoader
+    >>>
+    >>> config = EventConfig(enabled=True, device="cuda")
+    >>> detector = LattifAIEventDetector(config)
+    >>>
+    >>> audio = AudioLoader.load("speech.wav")
+    >>> result = detector.detect(audio)
+    >>>
+    >>> # Access VAD segments directly
+    >>> for start, end in result.vad_segments:
+    ...     print(f"Speech: {start:.2f} - {end:.2f}")
+    >>>
+    >>> # Or access the full TextGrid
+    >>> print(result.audio_events)
+
+Performance Notes:
+    - GPU acceleration provides significant speedup (10x+ over CPU)
+    - Use dtype="float16" for faster inference with minimal accuracy loss
+    - fast_mode=True reduces computation by only detecting top_k classes
+    - Long audio files are automatically chunked to manage memory
+
+See Also:
+    - lattifai.config.EventConfig: Configuration options
+    - lattifai_core.event: Core event detection implementation
+"""
+
+from .lattifai import LattifAIEventDetector
+
+__all__ = ["LattifAIEventDetector"]