lattifai 0.1.5__py3-none-any.whl → 0.2.2__py3-none-any.whl

lattifai/__init__.py

@@ -22,54 +22,19 @@ def _check_and_install_k2():
     """Check if k2 is installed and attempt to install it if not."""
     try:
         import k2
-
-        return True
     except ImportError:
-        pass
-
-    # k2 not found, try to install it
-    if os.environ.get('SKIP_K2_INSTALL'):
-        warnings.warn(
-            '\n' + '=' * 70 + '\n'
-            '  k2 is not installed and auto-installation is disabled.\n'
-            '  \n'
-            '  To use lattifai, please install k2 by running:\n'
-            '  \n'
-            '      install-k2\n'
-            '  \n' + '=' * 70,
-            RuntimeWarning,
-            stacklevel=2,
-        )
-        return False
-
-    print('\n' + '=' * 70)
-    print('  k2 is not installed. Attempting to install it now...')
-    print('  This is a one-time setup and may take a few minutes.')
-    print('=' * 70 + '\n')
-
-    try:
-        # Import and run the installation script
-        from scripts.install_k2 import install_k2_main
-
-        install_k2_main(dry_run=False)
-
-        print('\n' + '=' * 70)
-        print('  k2 has been installed successfully!')
-        print('=' * 70 + '\n')
-        return True
-    except Exception as e:
-        warnings.warn(
-            '\n' + '=' * 70 + '\n'
-            f'  Failed to auto-install k2: {e}\n'
-            '  \n'
-            '  Please install k2 manually by running:\n'
-            '  \n'
-            '      install-k2\n'
-            '  \n' + '=' * 70,
-            RuntimeWarning,
-            stacklevel=2,
-        )
-        return False
+        import subprocess
+
+        print('k2 is not installed. Attempting to install k2...')
+        try:
+            subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'install-k2'])
+            subprocess.check_call([sys.executable, '-m', 'install_k2'])
+            import k2  # Try importing again after installation
+
+            print('k2 installed successfully.')
+        except Exception as e:
+            warnings.warn(f'Failed to install k2 automatically. Please install it manually. Error: {e}')
+    return True
 
 
 # Auto-install k2 on first import

lattifai/bin/align.py

@@ -13,6 +13,23 @@ from lattifai.bin.cli_base import cli
     default='auto',
     help='Input Subtitle format.',
 )
+@click.option(
+    '-D',
+    '--device',
+    type=click.Choice(['cpu', 'cuda', 'mps'], case_sensitive=False),
+    default='cpu',
+    help='Device to use for inference.',
+)
+@click.option(
+    '-M', '--model_name_or_path', type=str, default='Lattifai/Lattice-1-Alpha', help='Lattifai model name or path'
+)
+@click.option(
+    '-S',
+    '--split_sentence',
+    is_flag=True,
+    default=False,
+    help='Re-segment subtitles by semantics.',
+)
 @click.argument(
     'input_audio_path',
     type=click.Path(exists=True, dir_okay=False),
@@ -30,13 +47,20 @@ def align(
     input_subtitle_path: Pathlike,
     output_subtitle_path: Pathlike,
     input_format: str = 'auto',
+    device: str = 'cpu',
+    model_name_or_path: str = 'Lattifai/Lattice-1-Alpha',
+    split_sentence: bool = False,
 ):
     """
     Command used to align audio with subtitles
     """
     from lattifai import LattifAI
 
-    client = LattifAI()
+    client = LattifAI(model_name_or_path=model_name_or_path, device=device)
     client.alignment(
-        input_audio_path, input_subtitle_path, format=input_format, output_subtitle_path=output_subtitle_path
+        input_audio_path,
+        input_subtitle_path,
+        format=input_format.lower(),
+        split_sentence=split_sentence,
+        output_subtitle_path=output_subtitle_path,
     )

lattifai/bin/cli_base.py

@@ -8,6 +8,11 @@ def cli():
     """
     The shell entry point to Lattifai, a tool for audio data manipulation.
     """
+    # Load environment variables from .env file
+    from dotenv import load_dotenv
+
+    load_dotenv()
+
     logging.basicConfig(
         format='%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s',
         level=logging.INFO,

lattifai/client.py

@@ -11,7 +11,7 @@ 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.tokenizer import LatticeTokenizer
 from lattifai.workers import Lattice1AlphaWorker
 
 load_dotenv()
@@ -25,9 +25,9 @@ class LattifAI(SyncAPIClient):
         *,
         api_key: Optional[str] = None,
         model_name_or_path: str = 'Lattifai/Lattice-1-Alpha',
-        device: str = 'cpu',
+        device: Optional[str] = None,
         base_url: Optional[str] = None,
-        timeout: Union[float, int] = 60.0,
+        timeout: Union[float, int] = 120.0,
         max_retries: int = 2,
         default_headers: Optional[Dict[str, str]] = None,
     ) -> None:
@@ -55,11 +55,26 @@ class LattifAI(SyncAPIClient):
         # Initialize components
         if not Path(model_name_or_path).exists():
             from huggingface_hub import snapshot_download
+            from huggingface_hub.errors import LocalEntryNotFoundError
 
-            model_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
+            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,
@@ -72,6 +87,7 @@ class LattifAI(SyncAPIClient):
         audio: Pathlike,
         subtitle: Pathlike,
         format: Optional[SubtitleFormat] = None,
+        split_sentence: bool = False,
         output_subtitle_path: Optional[Pathlike] = None,
     ) -> str:
         """Perform alignment on audio and subtitle/text.
@@ -87,11 +103,11 @@ class LattifAI(SyncAPIClient):
         # 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'))
+        print(colorful.green(f'         ✓ Parsed {len(supervisions)} subtitle 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)
+        lattice_id, lattice_graph = self.tokenizer.tokenize(supervisions, split_sentence=split_sentence)
         print(colorful.green(f'         ✓ Generated lattice graph with ID: {lattice_id}'))
 
         # step3: align audio with text
@@ -117,13 +133,10 @@ if __name__ == '__main__':
     import sys
 
     if len(sys.argv) == 4:
-        pass
+        audio, subtitle, output = sys.argv[1:]
     else:
         audio = 'tests/data/SA1.wav'
-        text = 'tests/data/SA1.TXT'
-
-    alignments = client.alignment(audio, text)
-    print(alignments)
+        subtitle = 'tests/data/SA1.TXT'
+        output = None
 
-    alignments = client.alignment(audio, 'not paired texttttt', format='txt')
-    print(alignments)
+    alignments = client.alignment(audio, subtitle, output_subtitle_path=output, split_sentence=True)

lattifai/io/reader.py

@@ -58,13 +58,12 @@ class SubtitleReader(ABCMeta):
                 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,
+                    # "duration": (event.end - event.start) / 1000.0 if event.end is not None else None,
                     # }
                 )
             )

lattifai/tokenizer/tokenizer.py

@@ -0,0 +1,284 @@
+import gzip
+import pickle
+import re
+from collections import defaultdict
+from itertools import chain
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from lattifai.base_client import SyncAPIClient
+from lattifai.io import Supervision
+from lattifai.tokenizer.phonemizer import G2Phonemizer
+
+PUNCTUATION = '!"#$%&()*+,-./:;<=>?@[\\]^_`{|}~'
+END_PUNCTUATION = '.!?"]。！？”】'
+PUNCTUATION_SPACE = PUNCTUATION + ' '
+STAR_TOKEN = '※'
+
+GROUPING_SEPARATOR = '✹'
+
+MAXIMUM_WORD_LENGTH = 40
+
+
+class LatticeTokenizer:
+    """Tokenizer for converting Lhotse Cut to LatticeGraph."""
+
+    def __init__(self, client_wrapper: SyncAPIClient):
+        self.client_wrapper = client_wrapper
+        self.words: List[str] = []
+        self.g2p_model: Any = None  # Placeholder for G2P model
+        self.dictionaries = defaultdict(lambda: [])
+        self.oov_word = '<unk>'
+        self.sentence_splitter = None
+        self.device = 'cpu'
+
+    def init_sentence_splitter(self):
+        if self.sentence_splitter is not None:
+            return
+
+        import onnxruntime as ort
+        from wtpsplit import SaT
+
+        providers = []
+        device = self.device
+        if device.startswith('cuda') and ort.get_all_providers().count('CUDAExecutionProvider') > 0:
+            providers.append('CUDAExecutionProvider')
+        elif device.startswith('mps') and ort.get_all_providers().count('MPSExecutionProvider') > 0:
+            providers.append('MPSExecutionProvider')
+
+        sat = SaT(
+            'sat-3l-sm',
+            ort_providers=providers + ['CPUExecutionProvider'],
+        )
+        self.sentence_splitter = sat
+
+    @staticmethod
+    def _resplit_special_sentence_types(sentence: str) -> List[str]:
+        """
+        Re-split special sentence types.
+
+        Examples:
+        '[APPLAUSE] &gt;&gt; MIRA MURATI:' -> ['[APPLAUSE]', '&gt;&gt; MIRA MURATI:']
+        '[MUSIC] &gt;&gt; SPEAKER:' -> ['[MUSIC]', '&gt;&gt; SPEAKER:']
+
+        Special handling patterns:
+        1. Separate special marks at the beginning (e.g., [APPLAUSE], [MUSIC], etc.) from subsequent speaker marks
+        2. Use speaker marks (&gt;&gt; or other separators) as split points
+
+        Args:
+            sentence: Input sentence string
+
+        Returns:
+            List of re-split sentences. If no special marks are found, returns the original sentence in a list
+        """
+        # Detect special mark patterns: [SOMETHING] &gt;&gt; SPEAKER:
+        # or other forms like [SOMETHING] SPEAKER:
+
+        # Pattern 1: [mark] HTML-encoded separator speaker:
+        pattern1 = r'^(\[[^\]]+\])\s+(&gt;&gt;|>>)\s+(.+)$'
+        match1 = re.match(pattern1, sentence.strip())
+        if match1:
+            special_mark = match1.group(1)
+            separator = match1.group(2)
+            speaker_part = match1.group(3)
+            return [special_mark, f'{separator} {speaker_part}']
+
+        # Pattern 2: [mark] speaker:
+        pattern2 = r'^(\[[^\]]+\])\s+([^:]+:)(.*)$'
+        match2 = re.match(pattern2, sentence.strip())
+        if match2:
+            special_mark = match2.group(1)
+            speaker_label = match2.group(2)
+            remaining = match2.group(3).strip()
+            if remaining:
+                return [special_mark, f'{speaker_label} {remaining}']
+            else:
+                return [special_mark, speaker_label]
+
+        # If no special pattern matches, return the original sentence
+        return [sentence]
+
+    @staticmethod
+    def from_pretrained(
+        client_wrapper: SyncAPIClient,
+        model_path: str,
+        device: str = 'cpu',
+        compressed: bool = True,
+    ):
+        """Load tokenizer from exported binary file"""
+        from pathlib import Path
+
+        words_model_path = f'{model_path}/words.bin'
+        if compressed:
+            with gzip.open(words_model_path, 'rb') as f:
+                data = pickle.load(f)
+        else:
+            with open(words_model_path, 'rb') as f:
+                data = pickle.load(f)
+
+        tokenizer = LatticeTokenizer(client_wrapper=client_wrapper)
+        tokenizer.words = data['words']
+        tokenizer.dictionaries = defaultdict(list, data['dictionaries'])
+        tokenizer.oov_word = data['oov_word']
+
+        g2p_model_path = f'{model_path}/g2p.bin' if Path(f'{model_path}/g2p.bin').exists() else None
+        if g2p_model_path:
+            tokenizer.g2p_model = G2Phonemizer(g2p_model_path, device=device)
+
+        tokenizer.device = device
+        tokenizer.add_special_tokens()
+        return tokenizer
+
+    def add_special_tokens(self):
+        tokenizer = self
+        for special_token in ['&gt;&gt;', '&gt;']:
+            if special_token not in tokenizer.dictionaries:
+                tokenizer.dictionaries[special_token] = tokenizer.dictionaries[tokenizer.oov_word]
+        return self
+
+    def prenormalize(self, texts: List[str], language: Optional[str] = None) -> List[str]:
+        if not self.g2p_model:
+            raise ValueError('G2P model is not loaded, cannot prenormalize texts')
+
+        oov_words = []
+        for text in texts:
+            words = text.lower().replace('-', ' ').replace('—', ' ').replace('–', ' ').split()
+            oovs = [w for w in words if w not in self.words]
+            if oovs:
+                oov_words.extend([w for w in oovs if (w not in self.words and len(w) <= MAXIMUM_WORD_LENGTH)])
+
+        oov_words = list(set(oov_words))
+        if oov_words:
+            indexs = []
+            for k, _word in enumerate(oov_words):
+                if any(_word.startswith(p) and _word.endswith(q) for (p, q) in [('(', ')'), ('[', ']')]):
+                    self.dictionaries[_word] = self.dictionaries[self.oov_word]
+                else:
+                    _word = _word.strip(PUNCTUATION_SPACE)
+                    if not _word or _word in self.words:
+                        indexs.append(k)
+            for idx in sorted(indexs, reverse=True):
+                del oov_words[idx]
+
+            g2p_words = [w for w in oov_words if w not in self.dictionaries]
+            if g2p_words:
+                predictions = self.g2p_model(words=g2p_words, lang=language, batch_size=len(g2p_words), num_prons=4)
+                for _word, _predictions in zip(g2p_words, predictions):
+                    for pronuncation in _predictions:
+                        if pronuncation and pronuncation not in self.dictionaries[_word]:
+                            self.dictionaries[_word].append(pronuncation)
+                    if not self.dictionaries[_word]:
+                        self.dictionaries[_word] = self.dictionaries[self.oov_word]
+
+            pronunciation_dictionaries: Dict[str, List[List[str]]] = {
+                w: self.dictionaries[w] for w in oov_words if self.dictionaries[w]
+            }
+            return pronunciation_dictionaries
+
+        return {}
+
+    def split_sentences(self, supervisions: List[Supervision], strip_whitespace=True) -> List[str]:
+        texts, text_len, sidx = [], 0, 0
+        for s, supervision in enumerate(supervisions):
+            text_len += len(supervision.text)
+            if text_len >= 2000 or s == len(supervisions) - 1:
+                text = ' '.join([sup.text for sup in supervisions[sidx : s + 1]])
+                texts.append(text)
+                sidx = s + 1
+                text_len = 0
+        if sidx < len(supervisions):
+            text = ' '.join([sup.text for sup in supervisions[sidx:]])
+            texts.append(text)
+        sentences = self.sentence_splitter.split(texts, threshold=0.15, strip_whitespace=strip_whitespace)
+
+        supervisions, remainder = [], ''
+        for _sentences in sentences:
+            # Process and re-split special sentence types
+            processed_sentences = []
+            for s, _sentence in enumerate(_sentences):
+                if remainder:
+                    _sentence = remainder + _sentence
+                    remainder = ''
+
+                # Detect and split special sentence types: e.g., '[APPLAUSE] &gt;&gt; MIRA MURATI:' -> ['[APPLAUSE]', '&gt;&gt; MIRA MURATI:']  # noqa: E501
+                resplit_parts = self._resplit_special_sentence_types(_sentence)
+                if any(resplit_parts[-1].endswith(sp) for sp in [':', '：']):
+                    if s < len(_sentences) - 1:
+                        _sentences[s + 1] = resplit_parts[-1] + ' ' + _sentences[s + 1]
+                    else:  # last part
+                        remainder = resplit_parts[-1] + ' ' + remainder
+                    processed_sentences.extend(resplit_parts[:-1])
+                else:
+                    processed_sentences.extend(resplit_parts)
+
+            _sentences = processed_sentences
+
+            if remainder:
+                _sentences[0] = remainder + _sentences[0]
+                remainder = ''
+
+            if any(_sentences[-1].endswith(ep) for ep in END_PUNCTUATION):
+                supervisions.extend(Supervision(text=s) for s in _sentences)
+            else:
+                supervisions.extend(Supervision(text=s) for s in _sentences[:-1])
+                remainder += _sentences[-1] + ' '
+
+        if remainder.strip():
+            supervisions.append(Supervision(text=remainder.strip()))
+
+        return supervisions
+
+    def tokenize(self, supervisions: List[Supervision], split_sentence: bool = False) -> Tuple[str, Dict[str, Any]]:
+        if split_sentence:
+            self.init_sentence_splitter()
+            supervisions = self.split_sentences(supervisions)
+
+        pronunciation_dictionaries = self.prenormalize([s.text for s in supervisions])
+        response = self.client_wrapper.post(
+            'tokenize',
+            json={
+                'supervisions': [s.to_dict() for s in supervisions],
+                'pronunciation_dictionaries': pronunciation_dictionaries,
+            },
+        )
+        if response.status_code != 200:
+            raise Exception(f'Failed to tokenize texts: {response.text}')
+        result = response.json()
+        lattice_id = result['id']
+        return lattice_id, (result['lattice_graph'], result['final_state'], result.get('acoustic_scale', 1.0))
+
+    def detokenize(
+        self,
+        lattice_id: str,
+        lattice_results: Tuple[torch.Tensor, Any, Any, float, float],
+        # return_supervisions: bool = True,
+        # return_details: bool = False,
+    ) -> List[Supervision]:
+        emission, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
+        response = self.client_wrapper.post(
+            'detokenize',
+            json={
+                'lattice_id': lattice_id,
+                'frame_shift': frame_shift,
+                'results': [t.to_dict() for t in results[0]],
+                'labels': labels[0],
+                'offset': offset,
+                'channel': channel,
+                'destroy_lattice': True,
+            },
+        )
+        if response.status_code != 200:
+            raise Exception(f'Failed to detokenize lattice: {response.text}')
+        result = response.json()
+        # if return_details:
+        #     raise NotImplementedError("return_details is not implemented yet")
+        return [Supervision.from_dict(s) for s in result['supervisions']]
+
+
+# Compute average score weighted by the span length
+def _score(spans):
+    if not spans:
+        return 0.0
+    # TokenSpan(token=token, start=start, end=end, score=scores[start:end].mean().item())
+    return round(sum(s.score * len(s) for s in spans) / sum(len(s) for s in spans), ndigits=4)

lattifai/workers/lattice1_alpha.py

@@ -17,8 +17,6 @@ class Lattice1AlphaWorker:
     """Worker for processing audio with LatticeGraph."""
 
     def __init__(self, model_path: Pathlike, device: str = 'cpu', num_threads: int = 8) -> None:
-        if device != 'cpu':
-            raise NotImplementedError(f'Only cpu is supported for now, got device={device}.')
         self.config = json.load(open(f'{model_path}/config.json'))
 
         # SessionOptions
@@ -29,8 +27,11 @@ class Lattice1AlphaWorker:
         sess_options.add_session_config_entry('session.intra_op.allow_spinning', '0')
 
         providers = []
-        if device.startswith('cuda') or ort.get_all_providers().count('CUDAExecutionProvider') > 0:
+        if device.startswith('cuda') and ort.get_all_providers().count('CUDAExecutionProvider') > 0:
             providers.append('CUDAExecutionProvider')
+        elif device.startswith('mps') and ort.get_all_providers().count('MPSExecutionProvider') > 0:
+            providers.append('MPSExecutionProvider')
+
         self.acoustic_ort = ort.InferenceSession(
             f'{model_path}/acoustic_opt.onnx',
             sess_options,
@@ -49,13 +50,29 @@ class Lattice1AlphaWorker:
         _start = time.time()
         # audio -> features -> emission
         features = self.extractor(audio)  # (1, T, D)
-        ort_inputs = {
-            'features': features.cpu().numpy(),
-            'feature_lengths': np.array([features.size(1)], dtype=np.int64),
-        }
-        emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+        if features.shape[1] > 6000:
+            features_list = torch.split(features, 6000, dim=1)
+            emissions = []
+            for features in features_list:
+                ort_inputs = {
+                    'features': features.cpu().numpy(),
+                    'feature_lengths': np.array([features.size(1)], dtype=np.int64),
+                }
+                emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+                emissions.append(emission)
+            emission = torch.cat(
+                [torch.from_numpy(emission).to(self.device) for emission in emissions], dim=1
+            )  # (1, T, vocab_size)
+        else:
+            ort_inputs = {
+                'features': features.cpu().numpy(),
+                'feature_lengths': np.array([features.size(1)], dtype=np.int64),
+            }
+            emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+            emission = torch.from_numpy(emission).to(self.device)
+
         self.timings['emission'] += time.time() - _start
-        return torch.from_numpy(emission).to(self.device)  # (1, T, vocab_size) torch
+        return emission  # (1, T, vocab_size) torch
 
     def load_audio(self, audio: Union[Pathlike, BinaryIO]) -> Tuple[torch.Tensor, int]:
         # load audio
@@ -104,9 +121,14 @@ class Lattice1AlphaWorker:
         self.timings['decoding_graph'] += time.time() - _start
 
         _start = time.time()
+        if self.device.type == 'mps':
+            device = 'cpu'  # k2 does not support mps yet
+        else:
+            device = self.device
+
         results, labels = align_segments(
-            emission.to(self.device) * acoustic_scale,
-            decoding_graph.to(self.device),
+            emission.to(device) * acoustic_scale,
+            decoding_graph.to(device),
             torch.tensor([emission.shape[1]], dtype=torch.int32),
             search_beam=100,
             output_beam=40,