lattifai 0.1.4__py3-none-any.whl

lattifai/__init__.py

@@ -0,0 +1,30 @@
+from .base_client import LattifAIError
+from .io import SubtitleIO
+
+try:
+    from importlib.metadata import version
+except ImportError:
+    # Python < 3.8
+    from importlib_metadata import version
+
+try:
+    __version__ = version('lattifai')
+except Exception:
+    __version__ = '0.1.0'  # fallback version
+
+
+# Lazy import for LattifAI to avoid dependency issues during basic import
+def __getattr__(name):
+    if name == 'LattifAI':
+        from .client import LattifAI
+
+        return LattifAI
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+__all__ = [
+    'LattifAI',  # noqa: F822
+    'LattifAIError',
+    'SubtitleIO',
+    '__version__',
+]

lattifai/base_client.py

@@ -0,0 +1,118 @@
+"""Base client classes for LattifAI SDK."""
+
+import os
+from abc import ABC
+from typing import Any, Awaitable, Callable, Dict, Optional, Union  # noqa: F401
+
+import httpx
+
+
+class LattifAIError(Exception):
+    """Base exception for LattifAI errors."""
+
+    pass
+
+
+class BaseAPIClient(ABC):
+    """Abstract base class for API clients."""
+
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: Union[float, httpx.Timeout] = 60.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 LattifAIError(
+                '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'
+            )
+
+        self._api_key = api_key
+        self._base_url = base_url
+        self._timeout = timeout
+        self._max_retries = max_retries
+
+        headers = {
+            'User-Agent': 'LattifAI/Python',
+            'Authorization': f'Bearer {self._api_key}',
+        }
+        if default_headers:
+            headers.update(default_headers)
+        self._default_headers = headers
+
+
+class SyncAPIClient(BaseAPIClient):
+    """Synchronous API client."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._client = httpx.Client(
+            base_url=self._base_url,
+            timeout=self._timeout,
+            headers=self._default_headers,
+        )
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def close(self) -> None:
+        """Close the HTTP client."""
+        self._client.close()
+
+    def _request(
+        self,
+        method: str,
+        url: str,
+        *,
+        json: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> httpx.Response:
+        """Make an HTTP request."""
+        return self._client.request(method=method, url=url, json=json, **kwargs)
+
+    def post(self, api_endpoint: str, *, json: Optional[Dict[str, Any]] = None, **kwargs) -> httpx.Response:
+        """Make a POST request to the specified API endpoint."""
+        return self._request('POST', api_endpoint, json=json, **kwargs)
+
+
+class AsyncAPIClient(BaseAPIClient):
+    """Asynchronous API client."""
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._client = httpx.AsyncClient(
+            base_url=self._base_url,
+            timeout=self._timeout,
+            headers=self._default_headers,
+        )
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self._client.aclose()
+
+    async def _request(
+        self,
+        method: str,
+        url: str,
+        *,
+        json: Optional[Dict[str, Any]] = None,
+        files: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> httpx.Response:
+        """Make an HTTP request."""
+        return await self._client.request(method=method, url=url, json=json, files=files, **kwargs)

lattifai/bin/__init__.py

@@ -0,0 +1,2 @@
+from .align import *  # noqa
+from .subtitle import *  # noqa

lattifai/bin/align.py

@@ -0,0 +1,42 @@
+import click
+import colorful
+from lhotse.utils import Pathlike
+
+from lattifai.bin.cli_base import cli
+
+
+@cli.command()
+@click.option(
+    '-F',
+    '--input_format',
+    type=click.Choice(['srt', 'vtt', 'ass', 'txt', 'auto'], case_sensitive=False),
+    default='auto',
+    help='Input Subtitle format.',
+)
+@click.argument(
+    'input_audio_path',
+    type=click.Path(exists=True, dir_okay=False),
+)
+@click.argument(
+    'input_subtitle_path',
+    type=click.Path(exists=True, dir_okay=False),
+)
+@click.argument(
+    'output_subtitle_path',
+    type=click.Path(allow_dash=True),
+)
+def align(
+    input_audio_path: Pathlike,
+    input_subtitle_path: Pathlike,
+    output_subtitle_path: Pathlike,
+    input_format: str = 'auto',
+):
+    """
+    Command used to align audio with subtitles
+    """
+    from lattifai import LattifAI
+
+    client = LattifAI()
+    client.alignment(
+        input_audio_path, input_subtitle_path, format=input_format, output_subtitle_path=output_subtitle_path
+    )

lattifai/bin/cli_base.py

@@ -0,0 +1,14 @@
+import logging
+
+import click
+
+
+@click.group()
+def cli():
+    """
+    The shell entry point to Lattifai, a tool for audio data manipulation.
+    """
+    logging.basicConfig(
+        format='%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s',
+        level=logging.INFO,
+    )

lattifai/bin/subtitle.py

@@ -0,0 +1,32 @@
+import click
+from lhotse.utils import Pathlike
+
+from lattifai.bin.cli_base import cli
+
+
+@cli.group()
+def subtitle():
+    """Group of commands used to convert subtitle format."""
+    pass
+
+
+@subtitle.command()
+@click.argument(
+    'input_subtitle_path',
+    type=click.Path(exists=True, dir_okay=False),
+)
+@click.argument(
+    'output_subtitle_path',
+    type=click.Path(allow_dash=True),
+)
+def convert(
+    input_subtitle_path: Pathlike,
+    output_subtitle_path: Pathlike,
+):
+    """
+    Convert subtitle file to another format.
+    """
+    import pysubs2
+
+    subtitle = pysubs2.load(input_subtitle_path)
+    subtitle.save(output_subtitle_path)

lattifai/client.py

@@ -0,0 +1,131 @@
+"""LattifAI client implementation."""
+
+import logging
+import os
+from pathlib import Path
+from typing import Any, Awaitable, BinaryIO, Callable, Dict, Optional, 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.tokenizers import LatticeTokenizer
+from lattifai.workers import Lattice1AlphaWorker
+
+load_dotenv()
+
+
+class LattifAI(SyncAPIClient):
+    """Synchronous LattifAI client."""
+
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        device: str = 'cpu',
+        timeout: Union[float, int] = 60.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 LattifAIError(
+                '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,
+        )
+
+        # Initialize components
+        model_name_or_path = '/Users/feiteng/GEEK/OmniCaptions/HF_models/Lattice-1-Alpha'
+
+        if not Path(model_name_or_path).exists():
+            from huggingface_hub import hf_hub_download
+
+            model_path = hf_hub_download(repo_id=model_name_or_path, repo_type='model')
+        else:
+            model_path = model_name_or_path
+
+        self.tokenizer = LatticeTokenizer.from_pretrained(
+            client_wrapper=self,
+            model_path=f'{model_path}/words.bin',
+            g2p_model_path=f'{model_path}/g2p.bin' if Path(f'{model_path}/g2p.bin').exists() else None,
+            device=device,
+        )
+        self.worker = Lattice1AlphaWorker(model_path, device=device, num_threads=8)
+
+    def alignment(
+        self,
+        audio: Pathlike,
+        subtitle: Pathlike,
+        format: Optional[SubtitleFormat] = None,
+        output_subtitle_path: Optional[Pathlike] = None,
+    ) -> str:
+        """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)
+
+        Returns:
+            Aligned subtitles in specified format
+        """
+        # 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)} supervision segments'))
+
+        # 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)
+        print(colorful.green(f'         ✓ Generated lattice graph with ID: {lattice_id}'))
+
+        # 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'))
+
+        # 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}'))
+
+        return output_subtitle_path or alignments
+
+
+if __name__ == '__main__':
+    client = LattifAI()
+    import sys
+
+    if len(sys.argv) == 4:
+        pass
+    else:
+        audio = 'tests/data/SA1.wav'
+        text = 'tests/data/SA1.TXT'
+
+    alignments = client.alignment(audio, text)
+    print(alignments)
+
+    alignments = client.alignment(audio, 'not paired texttttt', format='txt')
+    print(alignments)

lattifai/io/__init__.py

@@ -0,0 +1,22 @@
+from typing import List, Optional
+
+from lhotse.utils import Pathlike
+
+from .reader import SubtitleFormat, SubtitleReader
+from .supervision import Supervision
+from .writer import SubtitleWriter
+
+__all__ = ['SubtitleReader', 'SubtitleWriter', 'SubtitleIO', 'Supervision']
+
+
+class SubtitleIO:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def read(cls, subtitle: Pathlike, format: Optional[SubtitleFormat] = None) -> List[Supervision]:
+        return SubtitleReader.read(subtitle, format=format)
+
+    @classmethod
+    def write(cls, alignments: List[Supervision], output_path: Pathlike) -> Pathlike:
+        return SubtitleWriter.write(alignments, output_path)

lattifai/io/reader.py

@@ -0,0 +1,71 @@
+from abc import ABCMeta
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+from lhotse.utils import Pathlike
+
+from .supervision import Supervision
+
+SubtitleFormat = Literal['txt', 'srt', 'vtt', 'ass', 'auto']
+
+
+class SubtitleReader(ABCMeta):
+    """Parser for converting different subtitle formats to List[Supervision]."""
+
+    @classmethod
+    def read(cls, subtitle: Pathlike, format: Optional[SubtitleFormat] = None) -> List[Supervision]:
+        """Parse text and convert to Lhotse List[Supervision].
+
+        Args:
+            text: Input text to parse. Can be either:
+                - str: Direct text content to parse
+                - Path: File path to read and parse
+            format: Input text format (txt, srt, vtt, ass, textgrid)
+
+        Returns:
+            Parsed text in Lhotse Cut
+        """
+        if not format and Path(str(subtitle)).exists():
+            format = Path(str(subtitle)).suffix.lstrip('.').lower()
+        elif format:
+            format = format.lower()
+
+        if format == 'txt' or subtitle[-4:].lower() == '.txt':
+            if not Path(str(subtitle)).exists():  # str
+                lines = [line.strip() for line in subtitle.split('\n')]
+            else:  # file
+                lines = [line.strip() for line in open(subtitle).readlines()]
+            examples = [Supervision(text=line) for line in lines if line]
+        else:
+            examples = cls._parse_subtitle(subtitle, format=format)
+
+        return examples
+
+    @classmethod
+    def _parse_subtitle(cls, subtitle: Pathlike, format: Optional[SubtitleFormat]) -> List[Supervision]:
+        import pysubs2
+
+        try:
+            subs: pysubs2.SSAFile = pysubs2.load(
+                subtitle, encoding='utf-8', format_=format if format != 'auto' else None
+            )  # file
+        except IOError:
+            try:
+                subs: pysubs2.SSAFile = pysubs2.SSAFile.from_string(
+                    subtitle, format_=format if format != 'auto' else None
+                )  # str
+            except:
+                subs: pysubs2.SSAFile = pysubs2.load(subtitle, encoding='utf-8')  # auto detect format
+
+        supervisions = []
+
+        for event in subs.events:
+            supervisions.append(
+                Supervision(
+                    text=event.text,
+                    # "start": event.start / 1000.0 if event.start is not None else None,
+                    # "duration": event.end / 1000.0 if event.end is not None else None,
+                    # }
+                )
+            )
+        return supervisions

lattifai/io/supervision.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from lhotse.supervision import SupervisionSegment
+from lhotse.utils import Seconds
+
+
+@dataclass
+class Supervision(SupervisionSegment):
+    text: Optional[str] = None
+    id: str = ''
+    recording_id: str = ''
+    start: Seconds = 0.0
+    duration: Seconds = 0.0
+
+
+__all__ = ['Supervision']

lattifai/io/writer.py

@@ -0,0 +1,49 @@
+from abc import ABCMeta
+from typing import List
+
+from lhotse.utils import Pathlike
+
+from .reader import SubtitleFormat, Supervision
+
+
+class SubtitleWriter(ABCMeta):
+    """Class for writing subtitle files."""
+
+    @classmethod
+    def write(cls, alignments: List[Supervision], output_path: Pathlike) -> Pathlike:
+        if str(output_path)[-4:].lower() == '.txt':
+            with open(output_path, 'w', encoding='utf-8') as f:
+                for sup in alignments:
+                    f.write(f'{sup.text}\n')
+        elif str(output_path)[-5:].lower() == '.json':
+            with open(output_path, 'w', encoding='utf-8') as f:
+                import json
+
+                json.dump([sup.to_dict() for sup in alignments], f, ensure_ascii=False, indent=4)
+        elif str(output_path).endswith('.TextGrid') or str(output_path).endswith('.textgrid'):
+            from tgt import Interval, IntervalTier, TextGrid, write_to_file
+
+            tg = TextGrid()
+            supervisions, words = [], []
+            for supervision in sorted(alignments, key=lambda x: x.start):
+                supervisions.append(Interval(supervision.start, supervision.end, supervision.text or ''))
+                if supervision.alignment and 'word' in supervision.alignment:
+                    for alignment in supervision.alignment['word']:
+                        words.append(Interval(alignment.start, alignment.end, alignment.symbol))
+
+            tg.add_tier(IntervalTier(name='utterances', objects=supervisions))
+            if words:
+                tg.add_tier(IntervalTier(name='words', objects=words))
+            write_to_file(tg, output_path, format='long')
+        else:
+            import pysubs2
+
+            subs = pysubs2.SSAFile()
+            for sup in alignments:
+                start = int(sup.start * 1000)
+                end = int(sup.end * 1000)
+                text = sup.text or ''
+                subs.append(pysubs2.SSAEvent(start=start, end=end, text=text))
+            subs.save(output_path)
+
+        return output_path

lattifai/tokenizers/__init__.py

@@ -0,0 +1,3 @@
+from .tokenizer import LatticeTokenizer
+
+__all__ = ['LatticeTokenizer']

lattifai/tokenizers/phonemizer.py

@@ -0,0 +1,50 @@
+import re
+from typing import List, Optional, Union
+
+from dp.phonemizer import Phonemizer
+from num2words import num2words
+
+LANGUAGE = 'omni'
+
+
+class G2Phonemizer:
+    def __init__(self, model_checkpoint, device):
+        self.phonemizer = Phonemizer.from_checkpoint(model_checkpoint, device=device).predictor
+        self.pattern = re.compile(r'\d+')
+
+    def num2words(self, word, lang: str):
+        matches = self.pattern.findall(word)
+        for match in matches:
+            word_equivalent = num2words(int(match), lang=lang)
+            word = word.replace(match, word_equivalent)
+        return word
+
+    def remove_special_tokens(self, decoded: List[str]) -> List[str]:
+        return [d for d in decoded if d not in self.phonemizer.phoneme_tokenizer.special_tokens]
+
+    def __call__(
+        self, words: Union[str, List[str]], lang: Optional[StopIteration], batch_size: int = 0, num_prons: int = 1
+    ):
+        is_list = True
+        if not isinstance(words, list):
+            words = [words]
+            is_list = False
+
+        predictions = self.phonemizer(
+            [self.num2words(word.replace(' .', '.').replace('.', ' .'), lang=lang or 'en') for word in words],
+            lang=LANGUAGE,
+            batch_size=min(batch_size or len(words), 128),
+            num_prons=num_prons,
+        )
+        if num_prons > 1:
+            predictions = [
+                [self.remove_special_tokens(_prediction.phoneme_tokens) for _prediction in prediction]
+                for prediction in predictions
+            ]
+        else:
+            predictions = [self.remove_special_tokens(prediction.phoneme_tokens) for prediction in predictions]
+
+        if is_list:
+            return predictions
+
+        return predictions[0]