lattifai 0.2.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

lattifai/io/gemini_reader.py

@@ -0,0 +1,371 @@
+"""Reader for YouTube transcript files with speaker labels and timestamps."""
+
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from lhotse.utils import Pathlike
+
+from .supervision import Supervision
+
+
+@dataclass
+class GeminiSegment:
+    """Represents a segment in the Gemini transcript with metadata."""
+
+    text: str
+    timestamp: Optional[float] = None
+    speaker: Optional[str] = None
+    section: Optional[str] = None
+    segment_type: str = 'dialogue'  # 'dialogue', 'event', or 'section_header'
+    line_number: int = 0
+
+    @property
+    def start(self) -> float:
+        """Return start time in seconds."""
+        return self.timestamp if self.timestamp is not None else 0.0
+
+
+class GeminiReader:
+    """Parser for YouTube transcript format with speaker labels and timestamps."""
+
+    # Regex patterns for parsing (supports both [HH:MM:SS] and [MM:SS] formats)
+    TIMESTAMP_PATTERN = re.compile(r'\[(\d{1,2}):(\d{2}):(\d{2})\]|\[(\d{1,2}):(\d{2})\]')
+    SECTION_HEADER_PATTERN = re.compile(r'^##\s*\[(\d{1,2}):(\d{2}):(\d{2})\]\s*(.+)$')
+    SPEAKER_PATTERN = re.compile(r'^\*\*(.+?[:：])\*\*\s*(.+)$')
+    EVENT_PATTERN = re.compile(r'^\[([^\]]+)\]\s*\[(?:(\d{1,2}):(\d{2}):(\d{2})|(\d{1,2}):(\d{2}))\]$')
+    INLINE_TIMESTAMP_PATTERN = re.compile(r'^(.+?)\s*\[(?:(\d{1,2}):(\d{2}):(\d{2})|(\d{1,2}):(\d{2}))\]$')
+
+    # New patterns for YouTube link format: [[MM:SS](URL&t=seconds)]
+    YOUTUBE_SECTION_PATTERN = re.compile(r'^##\s*\[\[(\d{1,2}):(\d{2})\]\([^)]*&t=(\d+)\)\]\s*(.+)$')
+    YOUTUBE_INLINE_PATTERN = re.compile(r'^(.+?)\s*\[\[(\d{1,2}):(\d{2})\]\([^)]*&t=(\d+)\)\]$')
+
+    @classmethod
+    def parse_timestamp(cls, *args) -> float:
+        """Convert timestamp to seconds.
+
+        Supports both HH:MM:SS and MM:SS formats.
+        Args can be (hours, minutes, seconds) or (minutes, seconds).
+        Can also accept a single argument which is seconds.
+        """
+        if len(args) == 3:
+            # HH:MM:SS format
+            hours, minutes, seconds = args
+            return int(hours) * 3600 + int(minutes) * 60 + int(seconds)
+        elif len(args) == 2:
+            # MM:SS format
+            minutes, seconds = args
+            return int(minutes) * 60 + int(seconds)
+        elif len(args) == 1:
+            # Direct seconds (from YouTube &t= parameter)
+            return int(args[0])
+        else:
+            raise ValueError(f'Invalid timestamp args: {args}')
+
+    @classmethod
+    def read(
+        cls,
+        transcript_path: Pathlike,
+        include_events: bool = False,
+        include_sections: bool = False,
+    ) -> List[GeminiSegment]:
+        """Parse YouTube transcript file and return list of transcript segments.
+
+        Args:
+                transcript_path: Path to the transcript file
+                include_events: Whether to include event descriptions like [Applause]
+                include_sections: Whether to include section headers
+
+        Returns:
+                List of GeminiSegment objects with all metadata
+        """
+        transcript_path = Path(transcript_path).expanduser().resolve()
+        if not transcript_path.exists():
+            raise FileNotFoundError(f'Transcript file not found: {transcript_path}')
+
+        segments: List[GeminiSegment] = []
+        current_section = None
+        current_speaker = None
+
+        with open(transcript_path, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+
+        for line_num, line in enumerate(lines, start=1):
+            line = line.strip()
+            if not line:
+                continue
+
+            # Skip table of contents
+            if line.startswith('* ['):
+                continue
+            if line.startswith('## Table of Contents'):
+                continue
+
+            # Parse section headers
+            section_match = cls.SECTION_HEADER_PATTERN.match(line)
+            if section_match:
+                hours, minutes, seconds, section_title = section_match.groups()
+                timestamp = cls.parse_timestamp(hours, minutes, seconds)
+                current_section = section_title.strip()
+                if include_sections:
+                    segments.append(
+                        GeminiSegment(
+                            text=section_title.strip(),
+                            timestamp=timestamp,
+                            section=current_section,
+                            segment_type='section_header',
+                            line_number=line_num,
+                        )
+                    )
+                continue
+
+            # Parse YouTube format section headers: ## [[MM:SS](URL&t=seconds)] Title
+            youtube_section_match = cls.YOUTUBE_SECTION_PATTERN.match(line)
+            if youtube_section_match:
+                minutes, seconds, url_seconds, section_title = youtube_section_match.groups()
+                # Use the URL seconds for more accuracy
+                timestamp = cls.parse_timestamp(url_seconds)
+                current_section = section_title.strip()
+                if include_sections:
+                    segments.append(
+                        GeminiSegment(
+                            text=section_title.strip(),
+                            timestamp=timestamp,
+                            section=current_section,
+                            segment_type='section_header',
+                            line_number=line_num,
+                        )
+                    )
+                continue
+
+            # Parse event descriptions [event] [HH:MM:SS] or [MM:SS]
+            event_match = cls.EVENT_PATTERN.match(line)
+            if event_match:
+                groups = event_match.groups()
+                event_text = groups[0]
+                # Parse timestamp - can be HH:MM:SS (groups 1,2,3) or MM:SS (groups 4,5)
+                if groups[1] is not None:  # HH:MM:SS format
+                    timestamp = cls.parse_timestamp(groups[1], groups[2], groups[3])
+                elif groups[4] is not None:  # MM:SS format
+                    timestamp = cls.parse_timestamp(groups[4], groups[5])
+                else:
+                    timestamp = None
+
+                if include_events and timestamp is not None:
+                    segments.append(
+                        GeminiSegment(
+                            text=event_text.strip(),
+                            timestamp=timestamp,
+                            section=current_section,
+                            segment_type='event',
+                            line_number=line_num,
+                        )
+                    )
+                continue
+
+            # Parse speaker dialogue: **Speaker:** Text [HH:MM:SS] or [MM:SS]
+            speaker_match = cls.SPEAKER_PATTERN.match(line)
+            if speaker_match:
+                speaker, text_with_timestamp = speaker_match.groups()
+                current_speaker = speaker.strip()
+
+                # Extract timestamp from the end of the text
+                timestamp_match = cls.INLINE_TIMESTAMP_PATTERN.match(text_with_timestamp.strip())
+                youtube_match = cls.YOUTUBE_INLINE_PATTERN.match(text_with_timestamp.strip())
+
+                if timestamp_match:
+                    groups = timestamp_match.groups()
+                    text = groups[0]
+                    # Parse timestamp - can be HH:MM:SS (groups 1,2,3) or MM:SS (groups 4,5)
+                    if groups[1] is not None:  # HH:MM:SS format
+                        timestamp = cls.parse_timestamp(groups[1], groups[2], groups[3])
+                    elif groups[4] is not None:  # MM:SS format
+                        timestamp = cls.parse_timestamp(groups[4], groups[5])
+                    else:
+                        timestamp = None
+                elif youtube_match:
+                    groups = youtube_match.groups()
+                    text = groups[0]
+                    # Extract seconds from URL parameter
+                    url_seconds = groups[3]
+                    timestamp = cls.parse_timestamp(url_seconds)
+                else:
+                    text = text_with_timestamp.strip()
+                    timestamp = None
+
+                segments.append(
+                    GeminiSegment(
+                        text=text.strip(),
+                        timestamp=timestamp,
+                        speaker=current_speaker,
+                        section=current_section,
+                        segment_type='dialogue',
+                        line_number=line_num,
+                    )
+                )
+                current_speaker = None  # Reset speaker after use
+                continue
+
+            # Parse plain text with timestamp at the end
+            inline_match = cls.INLINE_TIMESTAMP_PATTERN.match(line)
+            youtube_inline_match = cls.YOUTUBE_INLINE_PATTERN.match(line)
+
+            if inline_match:
+                groups = inline_match.groups()
+                text = groups[0]
+                # Parse timestamp - can be HH:MM:SS (groups 1,2,3) or MM:SS (groups 4,5)
+                if groups[1] is not None:  # HH:MM:SS format
+                    timestamp = cls.parse_timestamp(groups[1], groups[2], groups[3])
+                elif groups[4] is not None:  # MM:SS format
+                    timestamp = cls.parse_timestamp(groups[4], groups[5])
+                else:
+                    timestamp = None
+
+                segments.append(
+                    GeminiSegment(
+                        text=text.strip(),
+                        timestamp=timestamp,
+                        speaker=current_speaker,
+                        section=current_section,
+                        segment_type='dialogue',
+                        line_number=line_num,
+                    )
+                )
+                continue
+            elif youtube_inline_match:
+                groups = youtube_inline_match.groups()
+                text = groups[0]
+                # Extract seconds from URL parameter
+                url_seconds = groups[3]
+                timestamp = cls.parse_timestamp(url_seconds)
+
+                segments.append(
+                    GeminiSegment(
+                        text=text.strip(),
+                        timestamp=timestamp,
+                        speaker=current_speaker,
+                        section=current_section,
+                        segment_type='dialogue',
+                        line_number=line_num,
+                    )
+                )
+                continue
+
+            # Skip markdown headers and other formatting
+            if line.startswith('#'):
+                continue
+
+        return segments
+
+    @classmethod
+    def extract_for_alignment(
+        cls,
+        transcript_path: Pathlike,
+        merge_consecutive: bool = False,
+        min_duration: float = 0.1,
+        merge_max_gap: float = 2.0,
+    ) -> List[Supervision]:
+        """Extract text segments for forced alignment.
+
+        This extracts only dialogue segments (not events or section headers)
+        and converts them to Supervision objects suitable for alignment.
+
+        Args:
+                transcript_path: Path to the transcript file
+                merge_consecutive: Whether to merge consecutive segments from same speaker
+                min_duration: Minimum duration for a segment
+                merge_max_gap: Maximum time gap (seconds) to merge consecutive segments
+
+        Returns:
+                List of Supervision objects ready for alignment
+        """
+        segments = cls.read(transcript_path, include_events=False, include_sections=False)
+
+        # Filter to only dialogue segments with timestamps
+        dialogue_segments = [s for s in segments if s.segment_type == 'dialogue' and s.timestamp is not None]
+
+        if not dialogue_segments:
+            raise ValueError(f'No dialogue segments with timestamps found in {transcript_path}')
+
+        # Sort by timestamp
+        dialogue_segments.sort(key=lambda x: x.timestamp)
+
+        # Convert to Supervision objects
+        supervisions: List[Supervision] = []
+
+        for i, segment in enumerate(dialogue_segments):
+            # Estimate duration based on next segment
+            if i < len(dialogue_segments) - 1:
+                duration = dialogue_segments[i + 1].timestamp - segment.timestamp
+            else:
+                # Last segment: estimate based on text length (rough heuristic)
+                words = len(segment.text.split())
+                duration = words * 0.3  # ~0.3 seconds per word
+
+            supervisions.append(
+                Supervision(
+                    text=segment.text,
+                    start=segment.timestamp,
+                    duration=max(duration, min_duration),
+                    id=f'segment_{i:05d}',
+                    speaker=segment.speaker,
+                )
+            )
+
+        # Optionally merge consecutive segments from same speaker
+        if merge_consecutive:
+            merged = []
+            current_speaker = None
+            current_texts = []
+            current_start = None
+            last_end_time = None
+
+            for i, (segment, sup) in enumerate(zip(dialogue_segments, supervisions)):
+                # Check if we should merge with previous segment
+                should_merge = False
+                if segment.speaker == current_speaker and current_start is not None:
+                    # Same speaker - check time gap
+                    time_gap = sup.start - last_end_time if last_end_time else 0
+                    if time_gap <= merge_max_gap:
+                        should_merge = True
+
+                if should_merge:
+                    # Same speaker within time threshold, accumulate
+                    current_texts.append(segment.text)
+                    last_end_time = sup.start + sup.duration
+                else:
+                    # Different speaker or gap too large, save previous segment
+                    if current_texts:
+                        merged_text = ' '.join(current_texts)
+                        merged.append(
+                            Supervision(
+                                text=merged_text,
+                                start=current_start,
+                                duration=last_end_time - current_start,
+                                id=f'merged_{len(merged):05d}',
+                            )
+                        )
+                    current_speaker = segment.speaker
+                    current_texts = [segment.text]
+                    current_start = sup.start
+                    last_end_time = sup.start + sup.duration
+
+            # Add final segment
+            if current_texts:
+                merged_text = ' '.join(current_texts)
+                merged.append(
+                    Supervision(
+                        text=merged_text,
+                        start=current_start,
+                        duration=last_end_time - current_start,
+                        id=f'merged_{len(merged):05d}',
+                    )
+                )
+
+            supervisions = merged
+
+        return supervisions
+
+
+__all__ = ['GeminiReader', 'GeminiSegment']

lattifai/io/gemini_writer.py

@@ -0,0 +1,173 @@
+"""Writer for YouTube transcript files with corrected timestamps from alignment."""
+
+import re
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from lhotse.utils import Pathlike
+
+from .gemini_reader import GeminiReader, GeminiSegment
+from .supervision import Supervision
+
+
+class GeminiWriter:
+    """Writer for updating YouTube transcript timestamps based on alignment results."""
+
+    @staticmethod
+    def format_timestamp(seconds: float) -> str:
+        """Convert seconds to [HH:MM:SS] format."""
+        hours = int(seconds // 3600)
+        minutes = int((seconds % 3600) // 60)
+        secs = int(seconds % 60)
+        return f'[{hours:02d}:{minutes:02d}:{secs:02d}]'
+
+    @classmethod
+    def update_timestamps(
+        cls,
+        original_transcript: Pathlike,
+        aligned_supervisions: List[Supervision],
+        output_path: Pathlike,
+        timestamp_mapping: Optional[Dict[int, float]] = None,
+    ) -> Pathlike:
+        """Update transcript file with corrected timestamps from alignment.
+
+        Args:
+                original_transcript: Path to the original transcript file
+                aligned_supervisions: List of aligned Supervision objects with corrected timestamps
+                output_path: Path to write the updated transcript
+                timestamp_mapping: Optional manual mapping from line_number to new timestamp
+
+        Returns:
+                Path to the output file
+        """
+        original_path = Path(original_transcript)
+        output_path = Path(output_path)
+
+        # Read original file
+        with open(original_path, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+
+        # Parse original segments to get line numbers
+        original_segments = GeminiReader.read(original_transcript, include_events=True, include_sections=True)
+
+        # Create mapping from line number to new timestamp
+        if timestamp_mapping is None:
+            timestamp_mapping = cls._create_timestamp_mapping(original_segments, aligned_supervisions)
+
+        # Update timestamps in lines
+        updated_lines = []
+        for line_num, line in enumerate(lines, start=1):
+            if line_num in timestamp_mapping:
+                new_timestamp = timestamp_mapping[line_num]
+                updated_line = cls._replace_timestamp(line, new_timestamp)
+                updated_lines.append(updated_line)
+            else:
+                updated_lines.append(line)
+
+        # Write updated content
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, 'w', encoding='utf-8') as f:
+            f.writelines(updated_lines)
+
+        return output_path
+
+    @classmethod
+    def _create_timestamp_mapping(
+        cls, original_segments: List[GeminiSegment], aligned_supervisions: List[Supervision]
+    ) -> Dict[int, float]:
+        """Create mapping from line numbers to new timestamps based on alignment.
+
+        This performs text matching between original segments and aligned supervisions
+        to determine which timestamps should be updated.
+        """
+        mapping = {}
+
+        # Create a simple text-based matching
+        dialogue_segments = [s for s in original_segments if s.segment_type == 'dialogue']
+
+        # Try to match based on text content
+        for aligned_sup in aligned_supervisions:
+            aligned_text = aligned_sup.text.strip()
+
+            # Find best matching original segment
+            best_match = None
+            best_score = 0
+
+            for orig_seg in dialogue_segments:
+                orig_text = orig_seg.text.strip()
+
+                # Simple text similarity (could be improved with fuzzy matching)
+                if aligned_text == orig_text:
+                    best_match = orig_seg
+                    best_score = 1.0
+                    break
+                elif aligned_text in orig_text or orig_text in aligned_text:
+                    score = min(len(aligned_text), len(orig_text)) / max(len(aligned_text), len(orig_text))
+                    if score > best_score:
+                        best_score = score
+                        best_match = orig_seg
+
+            # If we found a good match, update the mapping
+            if best_match and best_score > 0.8:
+                mapping[best_match.line_number] = aligned_sup.start
+
+        return mapping
+
+    @classmethod
+    def _replace_timestamp(cls, line: str, new_timestamp: float) -> str:
+        """Replace timestamp in a line with new timestamp."""
+        new_ts_str = cls.format_timestamp(new_timestamp)
+
+        # Replace timestamp patterns
+        # Pattern 1: [HH:MM:SS] at the end or in brackets
+        line = re.sub(r'\[\d{2}:\d{2}:\d{2}\]', new_ts_str, line)
+
+        return line
+
+    @classmethod
+    def write_aligned_transcript(
+        cls,
+        aligned_supervisions: List[Supervision],
+        output_path: Pathlike,
+        include_word_timestamps: bool = False,
+    ) -> Pathlike:
+        """Write a new transcript file from aligned supervisions.
+
+        This creates a simplified transcript format with accurate timestamps.
+
+        Args:
+                aligned_supervisions: List of aligned Supervision objects
+                output_path: Path to write the transcript
+                include_word_timestamps: Whether to include word-level timestamps if available
+
+        Returns:
+                Path to the output file
+        """
+        output_path = Path(output_path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(output_path, 'w', encoding='utf-8') as f:
+            f.write('# Aligned Transcript\n\n')
+
+            for i, sup in enumerate(aligned_supervisions):
+                # Write segment with timestamp
+                start_ts = cls.format_timestamp(sup.start)
+                f.write(f'{start_ts} {sup.text}\n')
+
+                # Optionally write word-level timestamps
+                if include_word_timestamps and hasattr(sup, 'alignment') and sup.alignment:
+                    if 'word' in sup.alignment:
+                        f.write('  Words: ')
+                        word_parts = []
+                        for word_info in sup.alignment['word']:
+                            word_ts = cls.format_timestamp(word_info['start'])
+                            word_parts.append(f'{word_info["symbol"]}{word_ts}')
+                        f.write(' '.join(word_parts))
+                        f.write('\n')
+
+                f.write('\n')
+
+        return output_path
+
+
+__all__ = ['GeminiWriter']

lattifai/io/reader.py

@@ -30,16 +30,29 @@ class SubtitleReader(ABCMeta):
         elif format:
             format = format.lower()
 
-        if format == 'txt' or subtitle[-4:].lower() == '.txt':
+        if format == 'gemini' or str(subtitle).endswith('Gemini.md'):
+            from .gemini_reader import GeminiReader
+
+            supervisions = GeminiReader.extract_for_alignment(subtitle)
+        elif format == 'txt' or (format == 'auto' and str(subtitle)[-4:].lower() == '.txt'):
             if not Path(str(subtitle)).exists():  # str
-                lines = [line.strip() for line in subtitle.split('\n')]
+                lines = [line.strip() for line in str(subtitle).split('\n')]
             else:  # file
-                lines = [line.strip() for line in open(subtitle).readlines()]
-            examples = [Supervision(text=line) for line in lines if line]
+                path_str = str(subtitle)
+                with open(path_str, encoding='utf-8') as f:
+                    lines = [line.strip() for line in f.readlines()]
+            supervisions = [Supervision(text=line) for line in lines if line]
         else:
-            examples = cls._parse_subtitle(subtitle, format=format)
+            try:
+                supervisions = cls._parse_subtitle(subtitle, format=format)
+            except Exception as e:
+                del e
+                print(f"Failed to parse subtitle with format {format}, trying 'gemini' parser.")
+                from .gemini_reader import GeminiReader
 
-        return examples
+                supervisions = GeminiReader.extract_for_alignment(subtitle)
+
+        return supervisions
 
     @classmethod
     def _parse_subtitle(cls, subtitle: Pathlike, format: Optional[SubtitleFormat]) -> List[Supervision]:
@@ -62,9 +75,8 @@ class SubtitleReader(ABCMeta):
             supervisions.append(
                 Supervision(
                     text=event.text,
-                    # "start": event.start / 1000.0 if event.start is not None else None,
-                    # "duration": (event.end - event.start) / 1000.0 if event.end is not None else None,
-                    # }
+                    start=event.start / 1000.0 if event.start is not None else None,
+                    duration=(event.end - event.start) / 1000.0 if event.end is not None else None,
                 )
             )
         return supervisions

lattifai/io/supervision.py

@@ -7,6 +7,22 @@ from lhotse.utils import Seconds
 
 @dataclass
 class Supervision(SupervisionSegment):
+    """
+    Extended SupervisionSegment with simplified initialization.
+
+    Note: The `alignment` field is inherited from SupervisionSegment:
+        alignment: Optional[Dict[str, List[AlignmentItem]]] = None
+
+    Structure of alignment when return_details=True:
+        {
+            'word': [
+                AlignmentItem(symbol='hello', start=0.0, duration=0.5, score=0.95),
+                AlignmentItem(symbol='world', start=0.6, duration=0.4, score=0.92),
+                ...
+            ]
+        }
+    """
+
     text: Optional[str] = None
     id: str = ''
     recording_id: str = ''

lattifai/io/utils.py

@@ -0,0 +1,15 @@
+"""
+Utility constants and helper functions for subtitle I/O operations
+"""
+
+# Supported subtitle formats for reading/writing
+SUBTITLE_FORMATS = ['srt', 'vtt', 'ass', 'ssa', 'sub', 'sbv', 'txt', 'md']
+
+# Input subtitle formats (includes special formats like 'auto' and 'gemini')
+INPUT_SUBTITLE_FORMATS = ['srt', 'vtt', 'ass', 'ssa', 'sub', 'sbv', 'txt', 'auto', 'gemini']
+
+# Output subtitle formats (includes special formats like 'TextGrid' and 'json')
+OUTPUT_SUBTITLE_FORMATS = ['srt', 'vtt', 'ass', 'ssa', 'sub', 'sbv', 'txt', 'TextGrid', 'json']
+
+# All subtitle formats combined (for file detection)
+ALL_SUBTITLE_FORMATS = list(set(SUBTITLE_FORMATS + ['TextGrid', 'json', 'gemini']))