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

lattifai/event/lattifai.py

@@ -0,0 +1,166 @@
+"""LattifAI Audio Event Detection implementation."""
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+from lattifai.audio2 import AudioData
+from lattifai.config.event import EventConfig
+from lattifai.logging import get_logger
+
+if TYPE_CHECKING:
+    from lattifai_core.event import LEDOutput
+
+    from lattifai.data import Caption
+
+
+formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
+logging.basicConfig(format=formatter, level=logging.INFO)
+
+
+class LattifAIEventDetector:
+    """
+    LattifAI Audio Event Detector.
+
+    This class provides a high-level interface for audio event detection,
+    wrapping the core LattifAIEventDetector from lattifai_core.
+
+    Attributes:
+        config: EventConfig configuration object.
+
+    Example:
+        >>> from lattifai.event import LattifAIEventDetector
+        >>> from lattifai.config import EventConfig
+        >>>
+        >>> config = EventConfig(enabled=True, device="cuda")
+        >>> detector = LattifAIEventDetector(config)
+        >>>
+        >>> # Detect events from audio data
+        >>> result = detector.detect(audio_data)
+        >>>
+        >>> # Access VAD segments directly
+        >>> for start, end in result.vad_segments:
+        ...     print(f"Speech: {start:.2f} - {end:.2f}")
+        >>>
+        >>> # Or access the full TextGrid
+        >>> for tier in result.audio_events.tiers:
+        ...     print(f"Event type: {tier.name}")
+    """
+
+    def __init__(self, config: EventConfig):
+        """
+        Initialize LattifAI Audio Event Detector.
+
+        Args:
+            config: EventConfig configuration.
+        """
+        self.config = config
+        self.logger = get_logger("event")
+        self._detector = None
+
+    @property
+    def name(self) -> str:
+        """Human-readable name of the detector."""
+        return "LattifAI_EventDetector"
+
+    @property
+    def detector(self):
+        """Lazy-load and return the audio event detector."""
+        if self._detector is None:
+            from lattifai_core.event import LattifAIEventDetector as CoreEventDetector
+
+            self._detector = CoreEventDetector.from_pretrained(
+                model_path=self.config.model_path,
+                device=self.config.device,
+                client_wrapper=self.config.client_wrapper,
+            )
+        return self._detector
+
+    def detect(
+        self,
+        input_media: AudioData,
+        vad_chunk_size: Optional[float] = None,
+        vad_max_gap: Optional[float] = None,
+        fast_mode: Optional[bool] = None,
+    ) -> "LEDOutput":
+        """
+        Detect audio events in the input audio.
+
+        Args:
+            input_media: Audio data to analyze.
+            vad_chunk_size: Override config vad_chunk_size.
+            vad_max_gap: Override config vad_max_gap.
+            fast_mode: Override config fast_mode.
+
+        Returns:
+            LEDOutput containing audio_events, event_names, vad_segments.
+        """
+        return self.detector(
+            audio=input_media,
+            vad_chunk_size=vad_chunk_size or self.config.vad_chunk_size,
+            vad_max_gap=vad_max_gap or self.config.vad_max_gap,
+            fast_mode=fast_mode if fast_mode is not None else self.config.fast_mode,
+            custom_aliases=self.config.event_aliases or {},
+        )
+
+    def profiling(self, reset: bool = False) -> str:
+        """Get profiling information for the detector."""
+        if self._detector is None:
+            return ""
+        return self.detector.profiling(reset=reset, logger=self.logger)
+
+    def detect_and_update_caption(
+        self,
+        caption: "Caption",
+        input_media: AudioData,
+        vad_chunk_size: Optional[float] = None,
+        vad_max_gap: Optional[float] = None,
+        fast_mode: Optional[bool] = None,
+    ) -> "Caption":
+        """
+        Run event detection and update caption with audio events.
+
+        This is the main entry point for integrating event detection with alignment.
+        When event_matching is enabled, it also updates caption timestamps for [Event] markers.
+
+        Args:
+            audio: AudioData to analyze
+            caption: Caption to update with event detection results
+
+        Returns:
+            Updated Caption with event field populated
+        """
+        # Event matching: update caption timestamps based on detected events
+        if self.config.event_matching:
+            # Get supervisions to process
+            supervisions = caption.alignments or caption.supervisions
+
+            led_output, supervisions = self.detector.detect_and_update_supervisions(
+                supervisions=supervisions,
+                audio=input_media,
+                vad_chunk_size=vad_chunk_size or self.config.vad_chunk_size,
+                vad_max_gap=vad_max_gap or self.config.vad_max_gap,
+                fast_mode=fast_mode if fast_mode is not None else self.config.fast_mode,
+                custom_aliases=self.config.event_aliases or {},
+                extra_events=self.config.extra_events or None,
+                time_tolerance=self.config.time_tolerance,
+                update_timestamps=self.config.update_timestamps,
+                duplicate_strategy=self.config.duplicate_strategy,
+            )
+            # Store LEDOutput in caption
+            caption.event = led_output
+
+            if caption.alignments:
+                caption.alignments = supervisions
+            else:
+                caption.supervisions = supervisions
+        else:
+            # Simple detection without event matching
+            led_output = self.detect(
+                input_media=input_media,
+                vad_chunk_size=vad_chunk_size,
+                vad_max_gap=vad_max_gap,
+                fast_mode=fast_mode,
+            )
+            caption.event = led_output
+
+        return caption

lattifai/mixin.py

@@ -8,12 +8,12 @@ import colorful
 from lhotse.utils import Pathlike
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption
+from lattifai.data import Caption
 from lattifai.errors import CaptionProcessingError
 from lattifai.utils import safe_print
 
 if TYPE_CHECKING:
-    from .config import AlignmentConfig, CaptionConfig, ClientConfig, DiarizationConfig, TranscriptionConfig
+    from .config import AlignmentConfig, DiarizationConfig, EventConfig, TranscriptionConfig
 
 
 class LattifAIClientMixin:
@@ -170,9 +170,10 @@ class LattifAIClientMixin:
         alignment_config: Optional["AlignmentConfig"],
         transcription_config: Optional["TranscriptionConfig"],
         diarization_config: Optional["DiarizationConfig"] = None,
+        event_config: Optional["EventConfig"] = None,
     ) -> tuple:
         """Initialize all configs with defaults if not provided."""
-        from .config import AlignmentConfig, DiarizationConfig, TranscriptionConfig
+        from .config import AlignmentConfig, DiarizationConfig, EventConfig, TranscriptionConfig
 
         if alignment_config is None:
             alignment_config = AlignmentConfig()
@@ -180,20 +181,24 @@ class LattifAIClientMixin:
             transcription_config = TranscriptionConfig()
         if diarization_config is None:
             diarization_config = DiarizationConfig()
+        if event_config is None:
+            event_config = EventConfig()
 
         from lattifai.utils import _resolve_model_path
 
-        if transcription_config is not None:
-            transcription_config.lattice_model_path = _resolve_model_path(
-                alignment_config.model_name, getattr(alignment_config, "model_hub", "huggingface")
-            )
+        model_path = _resolve_model_path(
+            alignment_config.model_name, getattr(alignment_config, "model_hub", "modelscope")
+        )
+        transcription_config.lattice_model_path = model_path
+        event_config.model_path = model_path
 
         # Set client_wrapper for all configs
         alignment_config.client_wrapper = self
         transcription_config.client_wrapper = self
         diarization_config.client_wrapper = self
+        event_config.client_wrapper = self
 
-        return alignment_config, transcription_config, diarization_config
+        return alignment_config, transcription_config, diarization_config, event_config
 
     def _init_shared_components(
         self,
@@ -284,18 +289,18 @@ class LattifAIClientMixin:
                 format=input_caption_format,
                 normalize_text=normalize_text if normalize_text is not None else self.caption_config.normalize_text,
             )
-            diarization_file = Path(str(input_caption)).with_suffix(".SpkDiar")
+            diarization_file = Path(str(input_caption)).with_suffix(".Diarization")
             if diarization_file.exists():
                 if verbose:
                     safe_print(colorful.cyan(f"📖 Step1b: Reading speaker diarization from {diarization_file}"))
-                caption.read_speaker_diarization(diarization_file)
-            events_file = Path(str(input_caption)).with_suffix(".AED")
-            if events_file.exists():
+                caption.read_diarization(diarization_file)
+            event_file = Path(str(input_caption)).with_suffix(".LED")
+            if event_file.exists():
                 if verbose:
-                    safe_print(colorful.cyan(f"📖 Step1c: Reading audio events from {events_file}"))
-                from tgt import read_textgrid
+                    safe_print(colorful.cyan(f"📖 Step1c: Reading audio events from {event_file}"))
+                from lattifai_core.event import LEDOutput
 
-                caption.audio_events = read_textgrid(events_file)
+                caption.event = LEDOutput.read(event_file)
 
             if verbose:
                 safe_print(colorful.green(f"         ✓ Parsed {len(caption)} caption segments"))
@@ -333,9 +338,9 @@ class LattifAIClientMixin:
                 karaoke_config=self.caption_config.karaoke,
             )
             diarization_file = Path(str(output_caption_path)).with_suffix(".SpkDiar")
-            if not diarization_file.exists() and caption.speaker_diarization:
+            if not diarization_file.exists() and caption.diarization:
                 safe_print(colorful.green(f"    Writing speaker diarization to: {diarization_file}"))
-                caption.write_speaker_diarization(diarization_file)
+                caption.write_diarization(diarization_file)
 
             safe_print(colorful.green(f"🎉🎉🎉🎉🎉 Caption file written to: {output_caption_path}"))
             return result

lattifai/transcription/base.py

@@ -7,8 +7,9 @@ from typing import List, Optional, Union
 import numpy as np
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption, Supervision
+from lattifai.caption import Supervision
 from lattifai.config import TranscriptionConfig
+from lattifai.data import Caption
 from lattifai.logging import get_logger
 
 

lattifai/transcription/gemini.py

@@ -11,6 +11,7 @@ from google.genai.types import GenerateContentConfig, Part, ThinkingConfig
 from lattifai.audio2 import AudioData
 from lattifai.caption import Supervision
 from lattifai.config import TranscriptionConfig
+from lattifai.data import Caption
 from lattifai.transcription.base import BaseTranscriber
 from lattifai.transcription.prompts import get_prompt_loader
 
@@ -245,18 +246,41 @@ class GeminiTranscriber(BaseTranscriber):
         return transcript
 
     def _get_transcription_prompt(self) -> str:
-        """Get (and cache) transcription system prompt from prompts module."""
+        """Get (and cache) transcription system prompt.
+
+        Priority:
+        1. Custom prompt from config.prompt (file path or text)
+        2. Default prompt from prompts/gemini/transcription_gem.txt
+        """
         if self._system_prompt is not None:
             return self._system_prompt
 
-        # Load prompt from prompts/gemini/transcription_gem.txt
-        prompt_loader = get_prompt_loader()
-        base_prompt = prompt_loader.get_gemini_transcription_prompt()
+        # Check for custom prompt
+        if self.config.prompt:
+            prompt_path = Path(self.config.prompt)
+            if prompt_path.exists() and prompt_path.is_file():
+                # Load from file
+                base_prompt = prompt_path.read_text(encoding="utf-8").strip()
+                if self.config.verbose:
+                    self.logger.info(f"📝 Using custom prompt from file: {prompt_path}")
+            else:
+                # Use as direct text
+                base_prompt = self.config.prompt
+                if self.config.verbose:
+                    self.logger.info("📝 Using custom prompt text")
+        else:
+            # Load default prompt from prompts/gemini/transcription_gem.txt
+            prompt_loader = get_prompt_loader()
+            base_prompt = prompt_loader.get_gemini_transcription_prompt()
 
         # Add language-specific instruction if configured
         if self.config.language:
             base_prompt += f"\n\n* Use {self.config.language} language for transcription."
 
+        # Add media description context if available
+        if self.config.description:
+            base_prompt += f"\n\n## Media Context\n\n{self.config.description}"
+
         self._system_prompt = base_prompt
         return self._system_prompt
 
@@ -287,14 +311,21 @@ class GeminiTranscriber(BaseTranscriber):
     def _get_generation_config(self) -> GenerateContentConfig:
         """Lazily build the generation config since it rarely changes."""
         if self._generation_config is None:
+            # Only include thinking_config if thinking mode is enabled
+            thinking_config = None
+            if self.config.thinking:
+                thinking_config = ThinkingConfig(
+                    include_thoughts=self.config.include_thoughts,
+                    thinking_budget=-1,
+                )
+
             self._generation_config = GenerateContentConfig(
                 system_instruction=self._get_transcription_prompt(),
                 response_modalities=["TEXT"],
-                thinking_config=ThinkingConfig(
-                    include_thoughts=False,
-                    thinking_budget=-1,
-                    # thinking_level="high",  # "low", "medium"
-                ),
+                thinking_config=thinking_config,
+                temperature=self.config.temperature,
+                top_k=self.config.top_k,
+                top_p=self.config.top_p,
             )
         return self._generation_config
 
@@ -323,23 +354,123 @@ class GeminiTranscriber(BaseTranscriber):
             ),
         )
 
-        if not response.text:
-            raise RuntimeError("Empty response from Gemini API")
-
-        transcript = response.text.strip()
+        # Extract content based on include_thoughts setting
+        if self.config.include_thoughts:
+            transcript = self._extract_with_thoughts(response)
+        else:
+            if not response.text:
+                raise RuntimeError("Empty response from Gemini API")
+            transcript = response.text.strip()
 
         if self.config.verbose:
             self.logger.info(f"✅ Transcription completed ({source}): {len(transcript)} characters")
 
         return transcript
 
+    def _extract_with_thoughts(self, response) -> str:
+        """Extract response content including thinking process and metadata."""
+        output_parts = []
+        thoughts = []
+        text_parts = []
+
+        # Iterate through all parts in the response
+        for candidate in response.candidates:
+            for part in candidate.content.parts:
+                if hasattr(part, "thought") and part.thought:
+                    # This is a thinking part
+                    if hasattr(part, "text") and part.text:
+                        thoughts.append(part.text)
+                elif hasattr(part, "text") and part.text:
+                    # This is a regular text part
+                    text_parts.append(part.text)
+
+        # Extract metadata
+        metadata_lines = self._extract_response_metadata(response)
+        if metadata_lines:
+            output_parts.append("---")
+            output_parts.extend(metadata_lines)
+            output_parts.append("---\n")
+
+        # Format output with thoughts section if present
+        if thoughts:
+            output_parts.append("<thinking>")
+            output_parts.extend(thoughts)
+            output_parts.append("</thinking>\n")
+
+        output_parts.extend(text_parts)
+
+        result = "\n".join(output_parts).strip()
+        if not result:
+            raise RuntimeError("Empty response from Gemini API")
+
+        return result
+
+    def _extract_response_metadata(self, response) -> list:
+        """Extract useful metadata from Gemini response as YAML frontmatter."""
+        lines = []
+
+        # Model version
+        if hasattr(response, "model_version") and response.model_version:
+            lines.append(f"model_version: {response.model_version}")
+
+        # Usage metadata (token counts)
+        if hasattr(response, "usage_metadata") and response.usage_metadata:
+            usage = response.usage_metadata
+            if hasattr(usage, "prompt_token_count"):
+                lines.append(f"prompt_tokens: {usage.prompt_token_count}")
+            if hasattr(usage, "candidates_token_count"):
+                lines.append(f"output_tokens: {usage.candidates_token_count}")
+            if hasattr(usage, "total_token_count"):
+                lines.append(f"total_tokens: {usage.total_token_count}")
+            # Thinking tokens if available
+            if hasattr(usage, "thoughts_token_count") and usage.thoughts_token_count:
+                lines.append(f"thinking_tokens: {usage.thoughts_token_count}")
+
+        # Candidate-level metadata
+        if response.candidates:
+            candidate = response.candidates[0]
+
+            # Finish reason
+            if hasattr(candidate, "finish_reason") and candidate.finish_reason:
+                lines.append(f"finish_reason: {candidate.finish_reason}")
+
+            # Average log probability (confidence indicator)
+            if hasattr(candidate, "avg_logprobs") and candidate.avg_logprobs is not None:
+                lines.append(f"avg_logprobs: {candidate.avg_logprobs:.4f}")
+
+            # Citation metadata
+            if hasattr(candidate, "citation_metadata") and candidate.citation_metadata:
+                citations = getattr(candidate.citation_metadata, "citations", [])
+                if citations:
+                    lines.append("citations:")
+                    for cite in citations:
+                        uri = getattr(cite, "uri", "")
+                        start = getattr(cite, "start_index", "")
+                        end = getattr(cite, "end_index", "")
+                        if uri:
+                            lines.append(f"  - uri: {uri}")
+                            if start or end:
+                                lines.append(f"    range: [{start}, {end}]")
+
+        return lines
+
     def write(
-        self, transcript: str, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
+        self, transcript: Union[str, Caption], output_file: Path, encoding: str = "utf-8", cache_event: bool = True
     ) -> Path:
         """
-        Persist transcript text to disk and return the file path.
+        Persist transcript to disk and return the file path.
+
+        Supports both raw string (from transcribe_file) and Caption object
+        (after conversion in mixin._transcribe).
         """
         if isinstance(output_file, str):
             output_file = Path(output_file)
-        output_file.write_text(transcript, encoding=encoding)
+
+        if isinstance(transcript, Caption):
+            # Caption object - use its write method with gemini format
+            transcript.write(output_file, output_format="gemini")
+        else:
+            # Raw string from transcription
+            output_file.write_text(transcript, encoding=encoding)
+
         return output_file

lattifai/transcription/lattifai.py

@@ -6,8 +6,9 @@ from typing import List, Optional, Union
 import numpy as np
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption, Supervision
+from lattifai.caption import Supervision
 from lattifai.config import TranscriptionConfig
+from lattifai.data import Caption
 from lattifai.transcription.base import BaseTranscriber
 
 
@@ -53,8 +54,8 @@ class LattifAITranscriber(BaseTranscriber):
 
     async def transcribe_file(self, media_file: Union[str, Path, AudioData], language: Optional[str] = None) -> Caption:
         transcriber = self._ensure_transcriber()
-        transcription, audio_events = transcriber.transcribe(media_file, language=language, num_workers=2)
-        return Caption.from_transcription_results(transcription=transcription, audio_events=audio_events)
+        transcription, event = transcriber.transcribe(media_file, language=language, num_workers=2)
+        return Caption.from_transcription_results(transcription=transcription, event=event)
 
     def transcribe_numpy(
         self,
@@ -77,9 +78,7 @@ class LattifAITranscriber(BaseTranscriber):
             audio, language=language, return_hypotheses=True, progress_bar=False, timestamps=True
         )[0]
 
-    def write(
-        self, transcript: Caption, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
-    ) -> Path:
+    def write(self, transcript: Caption, output_file: Path, encoding: str = "utf-8", cache_event: bool = True) -> Path:
         """
         Persist transcript text to disk and return the file path.
         """
@@ -87,10 +86,8 @@ class LattifAITranscriber(BaseTranscriber):
             output_file,
             include_speaker_in_text=False,
         )
-        if cache_audio_events and transcript.audio_events:
-            from tgt import write_to_file
-
-            events_file = output_file.with_suffix(".AED")
-            write_to_file(transcript.audio_events, events_file, format="long")
+        if cache_event and transcript.event:
+            events_file = output_file.with_suffix(".LED")
+            transcript.event.write(events_file)
 
         return output_file

lattifai/types.py

@@ -5,7 +5,7 @@ from typing import List, TypeAlias, Union
 
 from lhotse.utils import Pathlike
 
-from .caption import Supervision
+from lattifai.caption import Supervision
 
 # Path-like types
 PathLike: TypeAlias = Pathlike  # Re-export for convenience (str | Path)