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

lattifai/bin/align.py

@@ -23,6 +23,13 @@ from lattifai.bin.cli_base import cli
 @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),
@@ -42,6 +49,7 @@ def align(
     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
@@ -50,5 +58,9 @@ def align(
 
     client = LattifAI(model_name_or_path=model_name_or_path, device=device)
     client.alignment(
-        input_audio_path, input_subtitle_path, format=input_format.lower(), 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/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()
@@ -87,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.
@@ -102,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
@@ -138,4 +139,4 @@ if __name__ == '__main__':
         subtitle = 'tests/data/SA1.TXT'
         output = None
 
-    alignments = client.alignment(audio, subtitle, output_subtitle_path=output)
+    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-0.2.0.dist-info → lattifai-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lattifai
-Version: 0.2.0
+Version: 0.2.2
 Summary: Lattifai Python SDK: Seamless Integration with Lattifai's Speech and Video AI Services
 Author-email: Lattifai Technologies <tech@lattifai.com>
 Maintainer-email: Lattice <tech@lattifai.com>
@@ -60,6 +60,7 @@ Requires-Dist: tgt
 Requires-Dist: onnxruntime
 Requires-Dist: resampy
 Requires-Dist: g2p-phonemizer==0.1.1
+Requires-Dist: wtpsplit>=2.1.6
 Provides-Extra: numpy
 Requires-Dist: numpy; extra == "numpy"
 Provides-Extra: test
@@ -117,9 +118,37 @@ Usage: lattifai align [OPTIONS] INPUT_AUDIO_PATH INPUT_SUBTITLE_PATH OUTPUT_SUBT
 Options:
   -F, --input_format [srt|vtt|ass|txt|auto]  Input Subtitle format.
   -D, --device [cpu|cuda|mps]                Device to use for inference.
+  --split_sentence                           Smart sentence splitting based on punctuation semantics.
   --help                                     Show this message and exit.
 ```
 
+#### Understanding --split_sentence
+
+The `--split_sentence` option performs intelligent sentence re-splitting based on punctuation and semantic boundaries. This is especially useful when processing subtitles that combine multiple semantic units in a single segment, such as:
+
+- **Mixed content**: Non-speech elements (e.g., `[APPLAUSE]`, `[MUSIC]`) followed by actual dialogue
+- **Natural punctuation boundaries**: Colons, periods, and other punctuation marks that indicate semantic breaks
+- **Concatenated phrases**: Multiple distinct utterances joined together without proper separation
+
+**Example transformations**:
+```
+Input:  "[APPLAUSE] >> MIRA MURATI: Thank you all"
+Output: ["[APPLAUSE]", ">> MIRA MURATI: Thank you all"]
+
+Input:  "[MUSIC] Welcome back. Today we discuss AI."
+Output: ["[MUSIC]", "Welcome back.", "Today we discuss AI."]
+```
+
+This feature helps improve alignment accuracy by:
+1. Respecting punctuation-based semantic boundaries
+2. Separating distinct utterances for more precise timing
+3. Maintaining semantic context for each independent phrase
+
+**Usage**:
+```bash
+lattifai align --split_sentence audio.wav subtitle.srt output.srt
+```
+
 ### Python API
 
 ```python
@@ -136,6 +165,7 @@ client = LattifAI(
 result = client.alignment(
     audio="audio.wav",
     subtitle="subtitle.srt",
+    split_sentence=False,
     output_subtitle_path="output.srt"
 )
 ```
@@ -161,13 +191,21 @@ LattifAI(
 
 ```python
 client.alignment(
-    audio: str,                    # Path to audio file
-    subtitle: str,                 # Path to subtitle/text file
-    format: Optional[str] = None,  # 'srt', 'vtt', 'ass', 'txt' (auto-detect if None)
+    audio: str,                           # Path to audio file
+    subtitle: str,                        # Path to subtitle/text file
+    format: Optional[str] = None,         # 'srt', 'vtt', 'ass', 'txt' (auto-detect if None)
+    split_sentence: bool = False,         # Smart sentence splitting based on punctuation semantics
     output_subtitle_path: Optional[str] = None
 ) -> str
 ```
 
+**Parameters**:
+- `audio`: Path to the audio file to be aligned
+- `subtitle`: Path to the subtitle or text file
+- `format`: Subtitle format ('srt', 'vtt', 'ass', 'txt'). Auto-detected if None
+- `split_sentence`: Enable intelligent sentence re-splitting (default: False). Set to True when subtitles combine multiple semantic units (non-speech elements + dialogue, or multiple sentences) that would benefit from separate timing alignment
+- `output_subtitle_path`: Output path for aligned subtitle (optional)
+
 ## Examples
 
 ### Basic Text Alignment
@@ -178,6 +216,7 @@ client.alignment(
     audio="speech.wav",
     subtitle="transcript.txt",
     format="txt",
+    split_sentence=False,
     output_subtitle_path="output.srt"
 )
 ```

{lattifai-0.2.0.dist-info → lattifai-0.2.2.dist-info}/RECORD

@@ -1,22 +1,22 @@
 lattifai/__init__.py,sha256=JXUg0dT74UyAtKOjewRs9ijr5sl9SYsc6oU_WItY314,1497
 lattifai/base_client.py,sha256=ktFtATjL9pLSJUD-VqeJKA1FHkrsGHX7Uq_x00H7gO8,3322
-lattifai/client.py,sha256=NjHUqMZFoRxuyxdzBNEcn5kU57gJl31FSb6i9DDcoCw,5059
+lattifai/client.py,sha256=QXbdTuDA5Aap2udu4iig7CVxlgwOIrydpuLlVASs0aA,5145
 lattifai/bin/__init__.py,sha256=7YhmtEM8kbxJtz2-KIskvpLKBZAvkMSceVx8z4fkgQ4,61
-lattifai/bin/align.py,sha256=nw-wABsNyxC8zN3siiqgNi1Foou3cYhVzIAomuVrFOY,1436
+lattifai/bin/align.py,sha256=nQs901SDYmxyH2AXBtjgZGzrpwLaxANQRYP49Bd1AWo,1669
 lattifai/bin/cli_base.py,sha256=y535WXDRX8StloFn9icpfw7nQt0JxuWBIuPMnRxAYy8,392
 lattifai/bin/subtitle.py,sha256=bUWImAHpvyY59Vskqb5loQiD5ytQOxR8lTQRiQ4LyNA,647
 lattifai/io/__init__.py,sha256=vHWRN7MvAch-GUeFqqO-gM57SM-4YOpGUjIxFJdjfPA,671
-lattifai/io/reader.py,sha256=ErPnPMUvYQpjZ7Vd86EsHUkOcEfKdoI8iM3yKHRzSOQ,2576
+lattifai/io/reader.py,sha256=mtgxT5c_BiHbqqJvPE3nf7TIe_OcWgGu1zr6iXasfrk,2591
 lattifai/io/supervision.py,sha256=5UfSsgBhXoDU3-6drDtoD7y8HIiA4xRKZnbOKgeejwM,354
 lattifai/io/writer.py,sha256=1eAEFLlL8kricxRDPFBtVmeC4IiFyFnjbWXvw0VU-q4,2036
-lattifai/tokenizers/__init__.py,sha256=aqv44PDtq6g3oFFKW_l4HSR5ywT5W8eP1dHHywIvBfs,72
-lattifai/tokenizers/phonemizer.py,sha256=SfRi1KIMpmaao6OVmR1h_I_3QU-vrE6D5bh72Afg5XM,1759
-lattifai/tokenizers/tokenizer.py,sha256=u4lgS6-jN9cLuMNIojA4Swfsqb1EcyI7Bh_iw7tuL-s,5818
+lattifai/tokenizer/__init__.py,sha256=aqv44PDtq6g3oFFKW_l4HSR5ywT5W8eP1dHHywIvBfs,72
+lattifai/tokenizer/phonemizer.py,sha256=SfRi1KIMpmaao6OVmR1h_I_3QU-vrE6D5bh72Afg5XM,1759
+lattifai/tokenizer/tokenizer.py,sha256=Yuo0pLPQnF2uX0Fm5g8i5vtcADn7GeLpSqdGpMJgTww,11492
 lattifai/workers/__init__.py,sha256=s6YfkIq4FDIAzY9sPjRpXnJfszj2repqnMTqydRM5Zw,83
 lattifai/workers/lattice1_alpha.py,sha256=1VFo59EcygEctTHOhkcII8v3_mrj8JEJ8Fcaqk_7LVo,5762
-lattifai-0.2.0.dist-info/licenses/LICENSE,sha256=LNuoH5jpXXNKgjQ3XLwztFq8D3O7kZI-LSg81o4ym2M,1065
-lattifai-0.2.0.dist-info/METADATA,sha256=bXb6z5D1C-9YwHeycSFs8SAhUp8VNJbE9u-J9lvYMZ8,8997
-lattifai-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lattifai-0.2.0.dist-info/entry_points.txt,sha256=CwTI2NbJvF9msIHboAfTA99cmDr_HOWoODjS8R64JOw,131
-lattifai-0.2.0.dist-info/top_level.txt,sha256=-OVWZ68YYFcTN13ARkLasp2OUappe9wEVq-CKes7jM4,17
-lattifai-0.2.0.dist-info/RECORD,,
+lattifai-0.2.2.dist-info/licenses/LICENSE,sha256=LNuoH5jpXXNKgjQ3XLwztFq8D3O7kZI-LSg81o4ym2M,1065
+lattifai-0.2.2.dist-info/METADATA,sha256=4vmPOYKsIlvADiw0zUDQ2dbDpe-vOV-o5A0Hs1p7xfg,10971
+lattifai-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lattifai-0.2.2.dist-info/entry_points.txt,sha256=CwTI2NbJvF9msIHboAfTA99cmDr_HOWoODjS8R64JOw,131
+lattifai-0.2.2.dist-info/top_level.txt,sha256=-OVWZ68YYFcTN13ARkLasp2OUappe9wEVq-CKes7jM4,17
+lattifai-0.2.2.dist-info/RECORD,,

lattifai/tokenizers/tokenizer.py

@@ -1,147 +0,0 @@
-import gzip
-import pickle
-from collections import defaultdict
-from itertools import chain
-from typing import Any, Dict, List, Optional, Tuple
-
-import torch
-
-from lattifai.base_client import SyncAPIClient
-from lattifai.io import Supervision
-from lattifai.tokenizers.phonemizer import G2Phonemizer
-
-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>'
-
-    @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)
-        return tokenizer
-
-    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)
-
-            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 tokenize(self, supervisions: List[Supervision]) -> Tuple[str, Dict[str, Any]]:
-        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)