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

lattifai/client.py

@@ -1,18 +1,26 @@
 """LattifAI client implementation."""
 
-import logging
+import asyncio
 import os
-from pathlib import Path
-from typing import Any, Awaitable, BinaryIO, Callable, Dict, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import colorful
 from dotenv import load_dotenv
 from lhotse.utils import Pathlike
 
-from lattifai.base_client import AsyncAPIClient, LattifAIError, SyncAPIClient
-from lattifai.io import SubtitleFormat, SubtitleIO
-from lattifai.tokenizer import LatticeTokenizer
-from lattifai.workers import Lattice1AlphaWorker
+from lattifai.base_client import AsyncAPIClient, SyncAPIClient
+from lattifai.errors import (
+    AlignmentError,
+    ConfigurationError,
+    LatticeDecodingError,
+    LatticeEncodingError,
+    LattifAIError,
+    SubtitleProcessingError,
+    handle_exception,
+)
+from lattifai.io import SubtitleFormat, SubtitleIO, Supervision
+from lattifai.tokenizer import AsyncLatticeTokenizer
+from lattifai.utils import _load_tokenizer, _load_worker, _resolve_model_path, _select_device
 
 load_dotenv()
 
@@ -34,7 +42,7 @@ class LattifAI(SyncAPIClient):
         if api_key is None:
             api_key = os.environ.get('LATTIFAI_API_KEY')
         if api_key is None:
-            raise LattifAIError(
+            raise ConfigurationError(
                 'The api_key client option must be set either by passing api_key to the client '
                 'or by setting the LATTIFAI_API_KEY environment variable'
             )
@@ -52,35 +60,12 @@ class LattifAI(SyncAPIClient):
             default_headers=default_headers,
         )
 
-        # Initialize components
-        if not Path(model_name_or_path).exists():
-            from huggingface_hub import snapshot_download
-            from huggingface_hub.errors import LocalEntryNotFoundError
+        model_path = _resolve_model_path(model_name_or_path)
+        device = _select_device(device)
 
-            try:
-                model_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
-            except LocalEntryNotFoundError:
-                os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
-                model_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
-        else:
-            model_path = model_name_or_path
-
-        # device setup
-        if device is None:
-            import torch
-
-            device = 'cpu'
-            if torch.backends.mps.is_available():
-                device = 'mps'
-            elif torch.cuda.is_available():
-                device = 'cuda'
-
-        self.tokenizer = LatticeTokenizer.from_pretrained(
-            client_wrapper=self,
-            model_path=model_path,
-            device=device,
-        )
-        self.worker = Lattice1AlphaWorker(model_path, device=device, num_threads=8)
+        self.tokenizer = _load_tokenizer(self, model_path, device)
+        self.worker = _load_worker(model_path, device)
+        self.device = device
 
     def alignment(
         self,
@@ -88,55 +73,243 @@ class LattifAI(SyncAPIClient):
         subtitle: Pathlike,
         format: Optional[SubtitleFormat] = None,
         split_sentence: bool = False,
+        return_details: bool = False,
         output_subtitle_path: Optional[Pathlike] = None,
-    ) -> str:
+    ) -> Tuple[List[Supervision], Optional[Pathlike]]:
         """Perform alignment on audio and subtitle/text.
 
         Args:
             audio: Audio file path
             subtitle: Subtitle/Text to align with audio
-            export_format: Output format (srt, vtt, ass, txt)
+            format: Input subtitle format (srt, vtt, ass, txt). Auto-detected if None
+            split_sentence: Enable intelligent sentence re-splitting based on punctuation semantics
+            return_details: Return word-level alignment details in Supervision.alignment field
+            output_subtitle_path: Output path for aligned subtitle (optional)
 
         Returns:
-            Aligned subtitles in specified format
+            Tuple containing:
+                - List of aligned Supervision objects with timing information
+                - Output subtitle path (if output_subtitle_path was provided)
+
+        Raises:
+            SubtitleProcessingError: If subtitle file cannot be parsed
+            LatticeEncodingError: If lattice graph generation fails
+            AlignmentError: If audio alignment fails
+            LatticeDecodingError: If lattice decoding fails
         """
-        # step1: parse text or subtitles
-        print(colorful.cyan(f'📖 Step 1: Reading subtitle file from {subtitle}'))
-        supervisions = SubtitleIO.read(subtitle, format=format)
-        print(colorful.green(f'         ✓ Parsed {len(supervisions)} subtitle segments'))
+        try:
+            # step1: parse text or subtitles
+            print(colorful.cyan(f'📖 Step 1: Reading subtitle file from {subtitle}'))
+            try:
+                supervisions = SubtitleIO.read(subtitle, format=format)
+                print(colorful.green(f'         ✓ Parsed {len(supervisions)} subtitle segments'))
+            except Exception as e:
+                raise SubtitleProcessingError(
+                    f'Failed to parse subtitle file: {subtitle}',
+                    subtitle_path=str(subtitle),
+                    context={'original_error': str(e)},
+                )
+
+            # step2: make lattice by call Lattifai API
+            print(colorful.cyan('🔗 Step 2: Creating lattice graph from segments'))
+            try:
+                supervisions, lattice_id, lattice_graph = self.tokenizer.tokenize(
+                    supervisions, split_sentence=split_sentence
+                )
+                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)
+
+            # step3: search lattice graph with audio
+            print(colorful.cyan(f'🔍 Step 3: Searching lattice graph with audio: {audio}'))
+            try:
+                lattice_results = self.worker.alignment(audio, lattice_graph)
+                print(colorful.green('         ✓ Lattice search completed'))
+            except Exception as e:
+                raise AlignmentError(
+                    f'Audio alignment failed for {audio}',
+                    audio_path=str(audio),
+                    subtitle_path=str(subtitle),
+                    context={'original_error': str(e)},
+                )
 
-        # step2: make lattice by call Lattifai API
-        print(colorful.cyan('🔗 Step 2: Creating lattice graph from text'))
-        lattice_id, lattice_graph = self.tokenizer.tokenize(supervisions, split_sentence=split_sentence)
-        print(colorful.green(f'         ✓ Generated lattice graph with ID: {lattice_id}'))
+            # step4: decode lattice results to aligned segments
+            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
+                )
+                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)
+
+            # step5: export alignments to target format
+            if output_subtitle_path:
+                try:
+                    SubtitleIO.write(alignments, output_path=output_subtitle_path)
+                    print(colorful.green(f'🎉🎉🎉🎉🎉 Subtitle file written to: {output_subtitle_path}'))
+                except Exception as e:
+                    raise SubtitleProcessingError(
+                        f'Failed to write output file: {output_subtitle_path}',
+                        subtitle_path=str(output_subtitle_path),
+                        context={'original_error': str(e)},
+                    )
+            return (alignments, output_subtitle_path)
 
-        # step3: align audio with text
-        print(colorful.cyan(f'🎵 Step 3: Performing alignment on audio file: {audio}'))
-        lattice_results = self.worker.alignment(audio, lattice_graph)
-        print(colorful.green('         ✓ Alignment completed successfully'))
+        except (SubtitleProcessingError, LatticeEncodingError, AlignmentError, LatticeDecodingError):
+            # Re-raise our specific errors as-is
+            raise
+        except Exception as e:
+            # Catch any unexpected errors and wrap them
+            raise AlignmentError(
+                'Unexpected error during alignment process',
+                audio_path=str(audio),
+                subtitle_path=str(subtitle),
+                context={'original_error': str(e), 'error_type': e.__class__.__name__},
+            )
 
-        # step4: decode the lattice paths
-        print(colorful.cyan('🔍 Step 4: Decoding lattice paths to final alignments'))
-        alignments = self.tokenizer.detokenize(lattice_id, lattice_results)
-        print(colorful.green(f'         ✓ Decoded {len(alignments)} aligned segments'))
 
-        # step5: export alignments to target format
-        if output_subtitle_path:
-            SubtitleIO.write(alignments, output_path=output_subtitle_path)
-            print(colorful.green(f'🎉🎉🎉🎉🎉 Subtitle file written to: {output_subtitle_path}'))
+class AsyncLattifAI(AsyncAPIClient):
+    """Asynchronous LattifAI client."""
 
-        return output_subtitle_path or alignments
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        model_name_or_path: str = 'Lattifai/Lattice-1-Alpha',
+        device: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: Union[float, int] = 120.0,
+        max_retries: int = 2,
+        default_headers: Optional[Dict[str, str]] = None,
+    ) -> None:
+        if api_key is None:
+            api_key = os.environ.get('LATTIFAI_API_KEY')
+        if api_key is None:
+            raise ConfigurationError(
+                'The api_key client option must be set either by passing api_key to the client '
+                'or by setting the LATTIFAI_API_KEY environment variable'
+            )
+
+        if base_url is None:
+            base_url = os.environ.get('LATTIFAI_BASE_URL')
+        if not base_url:
+            base_url = 'https://api.lattifai.com/v1'
+
+        super().__init__(
+            api_key=api_key,
+            base_url=base_url,
+            timeout=timeout,
+            max_retries=max_retries,
+            default_headers=default_headers,
+        )
+
+        model_path = _resolve_model_path(model_name_or_path)
+        device = _select_device(device)
+
+        self.tokenizer = _load_tokenizer(self, model_path, device, tokenizer_cls=AsyncLatticeTokenizer)
+        self.worker = _load_worker(model_path, device)
+        self.device = device
+
+    async def alignment(
+        self,
+        audio: Pathlike,
+        subtitle: Pathlike,
+        format: Optional[SubtitleFormat] = None,
+        split_sentence: bool = False,
+        return_details: bool = False,
+        output_subtitle_path: Optional[Pathlike] = None,
+    ) -> Tuple[List[Supervision], Optional[Pathlike]]:
+        try:
+            print(colorful.cyan(f'📖 Step 1: Reading subtitle file from {subtitle}'))
+            try:
+                supervisions = await asyncio.to_thread(SubtitleIO.read, subtitle, format=format)
+                print(colorful.green(f'         ✓ Parsed {len(supervisions)} subtitle segments'))
+            except Exception as e:
+                raise SubtitleProcessingError(
+                    f'Failed to parse subtitle file: {subtitle}',
+                    subtitle_path=str(subtitle),
+                    context={'original_error': str(e)},
+                )
+
+            print(colorful.cyan('🔗 Step 2: Creating lattice graph from segments'))
+            try:
+                supervisions, lattice_id, lattice_graph = await self.tokenizer.tokenize(
+                    supervisions,
+                    split_sentence=split_sentence,
+                )
+                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)
+
+            print(colorful.cyan(f'🔍 Step 3: Searching lattice graph with audio: {audio}'))
+            try:
+                lattice_results = await asyncio.to_thread(self.worker.alignment, audio, lattice_graph)
+                print(colorful.green('         ✓ Lattice search completed'))
+            except Exception as e:
+                raise AlignmentError(
+                    f'Audio alignment failed for {audio}',
+                    audio_path=str(audio),
+                    subtitle_path=str(subtitle),
+                    context={'original_error': str(e)},
+                )
+
+            print(colorful.cyan('🎯 Step 4: Decoding lattice results to aligned segments'))
+            try:
+                alignments = await self.tokenizer.detokenize(
+                    lattice_id, lattice_results, supervisions=supervisions, return_details=return_details
+                )
+                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)
+
+            if output_subtitle_path:
+                try:
+                    await asyncio.to_thread(SubtitleIO.write, alignments, output_subtitle_path)
+                    print(colorful.green(f'🎉🎉🎉🎉🎉 Subtitle file written to: {output_subtitle_path}'))
+                except Exception as e:
+                    raise SubtitleProcessingError(
+                        f'Failed to write output file: {output_subtitle_path}',
+                        subtitle_path=str(output_subtitle_path),
+                        context={'original_error': str(e)},
+                    )
+
+            return (alignments, output_subtitle_path)
+
+        except (SubtitleProcessingError, LatticeEncodingError, AlignmentError, LatticeDecodingError):
+            raise
+        except Exception as e:
+            raise AlignmentError(
+                'Unexpected error during alignment process',
+                audio_path=str(audio),
+                subtitle_path=str(subtitle),
+                context={'original_error': str(e), 'error_type': e.__class__.__name__},
+            )
 
 
 if __name__ == '__main__':
     client = LattifAI()
     import sys
 
-    if len(sys.argv) == 4:
-        audio, subtitle, output = sys.argv[1:]
+    if len(sys.argv) == 5:
+        audio, subtitle, output, split_sentence = sys.argv[1:]
+        split_sentence = split_sentence.lower() in ('true', '1', 'yes')
     else:
         audio = 'tests/data/SA1.wav'
         subtitle = 'tests/data/SA1.TXT'
         output = None
+        split_sentence = False
 
-    alignments = client.alignment(audio, subtitle, output_subtitle_path=output, split_sentence=True)
+    (alignments, output_subtitle_path) = client.alignment(
+        audio, subtitle, output_subtitle_path=output, split_sentence=split_sentence, return_details=True
+    )

lattifai/errors.py

@@ -0,0 +1,257 @@
+"""Error handling and exception classes for LattifAI SDK."""
+
+import traceback
+from typing import Any, Dict, Optional
+
+import colorful
+
+# Error help messages
+LATTICE_DECODING_FAILURE_HELP = (
+    'Failed to decode lattice alignment. Possible reasons:\n\n'
+    '1) Audio and text content mismatch:\n'
+    '   - The transcript/subtitle does not accurately match the audio content\n'
+    '   - Text may be from a different version or section of the audio\n'
+    '   ⚠️  Note: Gemini transcription may occasionally skip large segments of audio, causing alignment failures.\n'
+    '       We will detect and fix this issue in the next version.\n\n'
+    '2) Unsupported audio type:\n'
+    '   - Singing is not yet supported, this will be optimized in future versions\n\n'
+    '💡 Troubleshooting tips:\n'
+    '   • Verify the transcript matches the audio by listening to a few segments\n'
+    '   • For YouTube videos, manually check if auto-generated transcript are accurate\n'
+    '       • Consider using a different transcription source if Gemini results are incomplete'
+)
+
+
+class LattifAIError(Exception):
+    """Base exception for LattifAI errors."""
+
+    def __init__(self, message: str, error_code: Optional[str] = None, context: Optional[Dict[str, Any]] = None):
+        """Initialize LattifAI error.
+
+        Args:
+            message: Error message
+            error_code: Optional error code for categorization
+            context: Optional context information about the error
+        """
+        super().__init__(message)
+        self.message = message
+        self.error_code = error_code or self.__class__.__name__
+        self.context = context or {}
+
+    def get_support_info(self) -> str:
+        """Get support information for users."""
+        return (
+            f'\n{colorful.green("🔧 Need help? Here are two ways to get support:")}\n'
+            f'   1. 📝 Create a GitHub issue: {colorful.green("https://github.com/lattifai/lattifai-python/issues")}\n'
+            '      Please include:\n'
+            '      - Your audio file format and duration\n'
+            "      - The text/subtitle content you're trying to align\n"
+            '      - This error message and stack trace\n'
+            f'   2. 💬 Join our Discord community: {colorful.green("https://discord.gg/vzmTzzZgNu")}\n'
+            '      Our team and community can help you troubleshoot\n'
+        )
+
+    def get_message(self) -> str:
+        """Return formatted error message without support information."""
+        base_message = f'{colorful.red(f"[{self.error_code}] {self.message}")}'
+        if self.context:
+            context_str = f'\n{colorful.yellow("Context:")} ' + ', '.join(f'{k}={v}' for k, v in self.context.items())
+            base_message += context_str
+        return base_message
+
+    def __str__(self) -> str:
+        """Return formatted error message without support information.
+
+        Note: Support info should be displayed explicitly at the CLI level,
+        not automatically appended to avoid duplication when errors are re-raised.
+        """
+        return self.get_message()
+
+
+class AudioProcessingError(LattifAIError):
+    """Error during audio processing operations."""
+
+    def __init__(self, message: str, audio_path: Optional[str] = None, **kwargs):
+        context = kwargs.get('context', {})
+        if audio_path:
+            context['audio_path'] = audio_path
+        kwargs['context'] = context
+        super().__init__(message, **kwargs)
+
+
+class AudioLoadError(AudioProcessingError):
+    """Error loading or reading audio file."""
+
+    def __init__(self, audio_path: str, original_error: Optional[Exception] = None, **kwargs):
+        message = f'Failed to load audio file: {colorful.red(audio_path)}'
+        if original_error:
+            message += f' - {colorful.red(str(original_error))}'
+
+        context = kwargs.get('context', {})
+        context.update({'audio_path': audio_path, 'original_error': str(original_error) if original_error else None})
+        kwargs['context'] = context
+
+        super().__init__(message, audio_path=audio_path, **kwargs)
+
+
+class AudioFormatError(AudioProcessingError):
+    """Error with audio format or codec."""
+
+    def __init__(self, audio_path: str, format_issue: str, **kwargs):
+        message = f'Audio format error for {colorful.red(audio_path)}: {colorful.red(format_issue)}'
+        context = kwargs.get('context', {})
+        context.update({'audio_path': audio_path, 'format_issue': format_issue})
+        kwargs['context'] = context
+        super().__init__(message, audio_path=audio_path, **kwargs)
+
+
+class SubtitleProcessingError(LattifAIError):
+    """Error during subtitle/text processing operations."""
+
+    def __init__(self, message: str, subtitle_path: Optional[str] = None, **kwargs):
+        context = kwargs.get('context', {})
+        if subtitle_path:
+            context['subtitle_path'] = subtitle_path
+        kwargs['context'] = context
+        super().__init__(message, **kwargs)
+
+
+class SubtitleParseError(SubtitleProcessingError):
+    """Error parsing subtitle or text file."""
+
+    def __init__(self, subtitle_path: str, parse_issue: str, **kwargs):
+        message = f'Failed to parse subtitle file {subtitle_path}: {parse_issue}'
+        context = kwargs.get('context', {})
+        context.update({'subtitle_path': subtitle_path, 'parse_issue': parse_issue})
+        kwargs['context'] = context
+        super().__init__(message, subtitle_path=subtitle_path, **kwargs)
+
+
+class AlignmentError(LattifAIError):
+    """Error during audio-text alignment process."""
+
+    def __init__(self, message: str, audio_path: Optional[str] = None, subtitle_path: Optional[str] = None, **kwargs):
+        context = kwargs.get('context', {})
+        if audio_path:
+            context['audio_path'] = audio_path
+        if subtitle_path:
+            context['subtitle_path'] = subtitle_path
+        kwargs['context'] = context
+        super().__init__(message, **kwargs)
+
+
+class LatticeEncodingError(AlignmentError):
+    """Error generating lattice graph from text."""
+
+    def __init__(self, text_content: str, original_error: Optional[Exception] = None, **kwargs):
+        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_content_length': len(text_content),
+                'text_preview': text_content[:100] + '...' if len(text_content) > 100 else text_content,
+                '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)}'
+
+        # 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))}'
+
+        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)
+
+    def get_message(self) -> str:
+        """Return formatted error message with help text."""
+        base_message = f'{colorful.red(f"[{self.error_code}]")} {self.message}'
+        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)}'
+        return base_message
+
+
+class ModelLoadError(LattifAIError):
+    """Error loading AI model."""
+
+    def __init__(self, model_name: str, original_error: Optional[Exception] = None, **kwargs):
+        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
+        super().__init__(message, **kwargs)
+
+
+class DependencyError(LattifAIError):
+    """Error with required dependencies."""
+
+    def __init__(self, dependency_name: str, install_command: Optional[str] = None, **kwargs):
+        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
+        super().__init__(message, **kwargs)
+
+
+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
+        super().__init__(message, **kwargs)
+
+
+class ConfigurationError(LattifAIError):
+    """Error with client configuration."""
+
+    def __init__(self, config_issue: str, **kwargs):
+        message = f'Configuration error: {config_issue}'
+        super().__init__(message, **kwargs)
+
+
+def handle_exception(func):
+    """Decorator to handle exceptions and convert them to LattifAI errors."""
+
+    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__,
+                'original_exception': e.__class__.__name__,
+                'traceback': traceback.format_exc(),
+            }
+            raise LattifAIError(error_msg, context=context) from e
+
+    return wrapper

lattifai/io/__init__.py

@@ -2,11 +2,31 @@ from typing import List, Optional
 
 from lhotse.utils import Pathlike
 
+from .gemini_reader import GeminiReader, GeminiSegment
+from .gemini_writer import GeminiWriter
 from .reader import SubtitleFormat, SubtitleReader
 from .supervision import Supervision
+from .utils import (
+    ALL_SUBTITLE_FORMATS,
+    INPUT_SUBTITLE_FORMATS,
+    OUTPUT_SUBTITLE_FORMATS,
+    SUBTITLE_FORMATS,
+)
 from .writer import SubtitleWriter
 
-__all__ = ['SubtitleReader', 'SubtitleWriter', 'SubtitleIO', 'Supervision']
+__all__ = [
+    'SubtitleReader',
+    'SubtitleWriter',
+    'SubtitleIO',
+    'Supervision',
+    'GeminiReader',
+    'GeminiWriter',
+    'GeminiSegment',
+    'SUBTITLE_FORMATS',
+    'INPUT_SUBTITLE_FORMATS',
+    'OUTPUT_SUBTITLE_FORMATS',
+    'ALL_SUBTITLE_FORMATS',
+]
 
 
 class SubtitleIO: