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

lattifai/transcription/gemini.py

@@ -0,0 +1,219 @@
+"""Gemini 2.5 Pro transcription module with config-driven architecture."""
+
+import asyncio
+from pathlib import Path
+from typing import Optional, Union
+
+from google import genai
+from google.genai.types import GenerateContentConfig, Part, ThinkingConfig
+
+from lattifai.audio2 import AudioData
+from lattifai.config import TranscriptionConfig
+from lattifai.transcription.base import BaseTranscriber
+from lattifai.transcription.prompts import get_prompt_loader
+
+
+class GeminiTranscriber(BaseTranscriber):
+    """
+    Gemini 2.5/3 Pro audio transcription with config-driven architecture.
+
+    Uses TranscriptionConfig for all behavioral settings.
+    """
+
+    # Transcriber metadata
+    file_suffix = ".md"
+
+    # The specific Gem URL
+    GEM_URL = "https://gemini.google.com/gem/1870ly7xvW2hU_umtv-LedGsjywT0sQiN"
+
+    def __init__(
+        self,
+        transcription_config: Optional[TranscriptionConfig] = None,
+    ):
+        """
+        Initialize Gemini transcriber.
+
+        Args:
+            transcription_config: Transcription configuration. If None, uses default.
+        """
+        super().__init__(config=transcription_config)
+
+        self._client: Optional[genai.Client] = None
+        self._generation_config: Optional[GenerateContentConfig] = None
+        self._system_prompt: Optional[str] = None
+
+        # Warn if API key not available
+        if not self.config.gemini_api_key:
+            self.logger.warning(
+                "⚠️ Gemini API key not provided. API key will be required when calling transcription methods."
+            )
+
+    @property
+    def name(self) -> str:
+        """Human-readable name of the transcriber."""
+        return f"{self.config.model_name}"
+
+    async def transcribe_url(self, url: str, language: Optional[str] = None) -> str:
+        """
+        Transcribe audio from URL using Gemini 2.5 Pro.
+
+        Args:
+            url: URL to transcribe (e.g., YouTube)
+            language: Optional language code for transcription (overrides config)
+
+        Returns:
+            Transcribed text
+
+        Raises:
+            ValueError: If API key not provided
+            RuntimeError: If transcription fails
+        """
+        if self.config.verbose:
+            self.logger.info(f"🎤 Starting Gemini transcription for: {url}")
+
+        try:
+            contents = Part.from_uri(file_uri=url, mime_type="video/*")
+            return await self._run_generation(contents, source=url)
+
+        except ImportError:
+            raise RuntimeError("Google GenAI SDK not installed. Please install with: pip install google-genai")
+        except Exception as e:
+            self.logger.error(f"Gemini transcription failed: {str(e)}")
+            raise RuntimeError(f"Gemini transcription failed: {str(e)}")
+
+    async def transcribe_file(self, media_file: Union[str, Path, AudioData], language: Optional[str] = None) -> str:
+        """
+        Transcribe audio/video from local file using Gemini 2.5 Pro.
+
+        Args:
+            media_file: Path to local audio/video file
+            language: Optional language code for transcription (overrides config)
+
+        Returns:
+            Transcribed text
+
+        Raises:
+            ValueError: If API key not provided
+            RuntimeError: If transcription fails
+        """
+        media_file = str(media_file)
+
+        if self.config.verbose:
+            self.logger.info(f"🎤 Starting Gemini transcription for file: {media_file}")
+
+        try:
+            client = self._get_client()
+
+            # Upload audio file
+            if self.config.verbose:
+                self.logger.info("📤 Uploading audio file to Gemini...")
+            media_file = client.files.upload(path=media_file)
+
+            contents = Part.from_uri(file_uri=media_file.uri, mime_type=media_file.mime_type)
+            return await self._run_generation(contents, source=media_file, client=client)
+
+        except ImportError:
+            raise RuntimeError("Google GenAI SDK not installed. Please install with: pip install google-genai")
+        except Exception as e:
+            self.logger.error(f"Gemini transcription failed: {str(e)}")
+            raise RuntimeError(f"Gemini transcription failed: {str(e)}")
+
+    def _get_transcription_prompt(self) -> str:
+        """Get (and cache) transcription system prompt from prompts module."""
+        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()
+
+        # Add language-specific instruction if configured
+        if self.config.language:
+            base_prompt += f"\n\n* Use {self.config.language} language for transcription."
+
+        self._system_prompt = base_prompt
+        return self._system_prompt
+
+    def get_gem_info(self) -> dict:
+        """Get information about the Gem being used."""
+        return {
+            "gem_name": "Media Transcription Gem",
+            "gem_url": self.GEM_URL,
+            "model": self.config.model_name,
+            "description": "Specialized Gem for media content transcription",
+        }
+
+    def _build_result(self, transcript: str, output_file: Path) -> dict:
+        """Augment the base result with Gemini-specific metadata."""
+        base_result = super()._build_result(transcript, output_file)
+        base_result.update({"model": self.config.model_name, "language": self.config.language})
+        return base_result
+
+    def _get_client(self) -> genai.Client:
+        """Lazily create the Gemini client when first needed."""
+        if not self.config.gemini_api_key:
+            raise ValueError("Gemini API key is required for transcription")
+
+        if self._client is None:
+            self._client = genai.Client(api_key=self.config.gemini_api_key)
+        return self._client
+
+    def _get_generation_config(self) -> GenerateContentConfig:
+        """Lazily build the generation config since it rarely changes."""
+        if self._generation_config is None:
+            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"
+                ),
+            )
+        return self._generation_config
+
+    async def _run_generation(
+        self,
+        contents: Part,
+        *,
+        source: str,
+        client: Optional[genai.Client] = None,
+    ) -> str:
+        """
+        Shared helper for sending generation requests and handling the response.
+        """
+        client = client or self._get_client()
+        config = self._get_generation_config()
+
+        if self.config.verbose:
+            self.logger.info(f"🔄 Sending transcription request to {self.config.model_name} ({source})...")
+
+        response = await asyncio.get_event_loop().run_in_executor(
+            None,
+            lambda: client.models.generate_content(
+                model=self.config.model_name,
+                contents=contents,
+                config=config,
+            ),
+        )
+
+        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 write(
+        self, transcript: str, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
+    ) -> Path:
+        """
+        Persist transcript text to disk and return the file path.
+        """
+        if isinstance(output_file, str):
+            output_file = Path(output_file)
+        output_file.write_text(transcript, encoding=encoding)
+        return output_file

lattifai/transcription/lattifai.py

@@ -0,0 +1,103 @@
+"""Transcription module with config-driven architecture."""
+
+from pathlib import Path
+from typing import Optional, Union
+
+from lattifai.audio2 import AudioData
+from lattifai.caption import Caption
+from lattifai.config import TranscriptionConfig
+from lattifai.transcription.base import BaseTranscriber
+from lattifai.transcription.prompts import get_prompt_loader  # noqa: F401
+
+
+class LattifAITranscriber(BaseTranscriber):
+    """
+    LattifAI local transcription with config-driven architecture.
+
+    Uses TranscriptionConfig for all behavioral settings.
+    Note: This transcriber only supports local file transcription, not URLs.
+    """
+
+    # Transcriber metadata
+    file_suffix = ".ass"
+    supports_url = False
+
+    def __init__(
+        self,
+        transcription_config: TranscriptionConfig,
+    ):
+        """
+        Initialize Gemini transcriber.
+
+        Args:
+            transcription_config: Transcription configuration. If None, uses default.
+        """
+        super().__init__(
+            config=transcription_config,
+        )
+
+        self._system_prompt: Optional[str] = None
+        self._transcriber = None
+
+    @property
+    def name(self) -> str:
+        return f"{self.config.model_name}"
+
+    async def transcribe_url(self, url: str, language: Optional[str] = None) -> str:
+        """
+        URL transcription not supported for LattifAI local models.
+
+        This method exists to satisfy the BaseTranscriber interface but
+        will never be called because supports_url = False and the base
+        class checks this flag before calling this method.
+
+        Args:
+            url: URL to transcribe (not supported)
+            language: Optional language code (not used)
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support URL transcription. "
+            f"Please download the file first and use transcribe_file()."
+        )
+
+    async def transcribe_file(self, media_file: Union[str, Path, AudioData], language: Optional[str] = None) -> Caption:
+        if self._transcriber is None:
+            from lattifai_core.transcription import LattifAITranscriber as CoreLattifAITranscriber
+
+            self._transcriber = CoreLattifAITranscriber.from_pretrained(model_config=self.config)
+
+        transcription, audio_events = self._transcriber.transcribe(media_file, language=language, num_workers=2)
+        caption = Caption.from_transcription_results(
+            transcription=transcription,
+            audio_events=audio_events,
+        )
+
+        return caption
+
+    def write(
+        self, transcript: Caption, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
+    ) -> Path:
+        """
+        Persist transcript text to disk and return the file path.
+        """
+        transcript.write(
+            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")
+
+        return output_file
+
+    def _get_transcription_prompt(self) -> str:
+        """Get (and cache) transcription system prompt from prompts module."""
+        if self._system_prompt is not None:
+            return self._system_prompt
+
+        base_prompt = ""  # TODO
+
+        self._system_prompt = base_prompt
+        return self._system_prompt

lattifai/types.py

@@ -0,0 +1,30 @@
+"""Common type definitions for LattifAI."""
+
+from pathlib import Path
+from typing import List, TypeAlias, Union
+
+from lhotse.utils import Pathlike
+
+from .caption import Supervision
+
+# Path-like types
+PathLike: TypeAlias = Pathlike  # Re-export for convenience (str | Path)
+
+# Caption types
+SupervisionList: TypeAlias = List[Supervision]
+"""List of caption segments with timing and text information."""
+
+# Media format types
+MediaFormat: TypeAlias = str
+"""Media format string (e.g., 'mp3', 'wav', 'mp4')."""
+
+# URL types
+URL: TypeAlias = str
+"""String representing a URL."""
+
+__all__ = [
+    "PathLike",
+    "SupervisionList",
+    "MediaFormat",
+    "URL",
+]

lattifai/utils.py

@@ -6,8 +6,6 @@ from pathlib import Path
 from typing import Any, Optional, Type
 
 from lattifai.errors import ModelLoadError
-from lattifai.tokenizer import LatticeTokenizer
-from lattifai.workers import Lattice1AlphaWorker
 
 
 def _get_cache_marker_path(cache_dir: Path) -> Path:
@@ -55,8 +53,8 @@ def _create_cache_marker(cache_dir: Path) -> None:
 
 def _resolve_model_path(model_name_or_path: str) -> str:
     """Resolve model path, downloading from Hugging Face when necessary."""
-    if Path(model_name_or_path).exists():
-        return model_name_or_path
+    if Path(model_name_or_path).expanduser().exists():
+        return str(Path(model_name_or_path).expanduser())
 
     from huggingface_hub import snapshot_download
     from huggingface_hub.constants import HF_HUB_CACHE
@@ -94,7 +92,7 @@ def _resolve_model_path(model_name_or_path: str) -> str:
 
 def _select_device(device: Optional[str]) -> str:
     """Select best available torch device when not explicitly provided."""
-    if device:
+    if device and device != "auto":
         return device
 
     import torch
@@ -105,29 +103,3 @@ def _select_device(device: Optional[str]) -> str:
     elif torch.cuda.is_available():
         detected = "cuda"
     return detected
-
-
-def _load_tokenizer(
-    client_wrapper: Any,
-    model_path: str,
-    device: str,
-    *,
-    tokenizer_cls: Type[LatticeTokenizer] = LatticeTokenizer,
-) -> LatticeTokenizer:
-    """Instantiate tokenizer with consistent error handling."""
-    try:
-        return tokenizer_cls.from_pretrained(
-            client_wrapper=client_wrapper,
-            model_path=model_path,
-            device=device,
-        )
-    except Exception as e:
-        raise ModelLoadError(f"tokenizer from {model_path}", original_error=e)
-
-
-def _load_worker(model_path: str, device: str) -> Lattice1AlphaWorker:
-    """Instantiate lattice worker with consistent error handling."""
-    try:
-        return Lattice1AlphaWorker(model_path, device=device, num_threads=8)
-    except Exception as e:
-        raise ModelLoadError(f"worker from {model_path}", original_error=e)

lattifai/workflow/__init__.py

@@ -0,0 +1,22 @@
+"""
+LattifAI Agentic Workflows
+
+This module provides agentic workflow capabilities for automated processing
+of multimedia content through intelligent agent-based pipelines.
+"""
+
+# Import transcript processing functionality
+
+
+from .base import WorkflowAgent, WorkflowResult, WorkflowStep
+from .file_manager import TRANSCRIBE_CHOICE, FileExistenceManager
+from .youtube import YouTubeDownloader
+
+__all__ = [
+    "WorkflowAgent",
+    "WorkflowStep",
+    "WorkflowResult",
+    "FileExistenceManager",
+    "YouTubeDownloader",
+    "TRANSCRIBE_CHOICE",
+]

lattifai/workflow/agents.py

@@ -0,0 +1,6 @@
+"""
+Caption Agents
+
+"""
+
+__all__ = []