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

lattifai/io/writer.py

@@ -1,49 +1,90 @@
+import json
 from abc import ABCMeta
-from typing import List
+from typing import Any, List, Optional
 
+import pysubs2
+from lhotse.supervision import AlignmentItem
 from lhotse.utils import Pathlike
 
-from .reader import SubtitleFormat, Supervision
+from .reader import Supervision
 
 
 class SubtitleWriter(ABCMeta):
-    """Class for writing subtitle files."""
+    """Class for writing subtitle files with optional word-level alignment."""
 
     @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')
+                    word_items = parse_alignment_from_supervision(sup)
+                    if word_items:
+                        for item in word_items:
+                            f.write(f'[{item.start:.2f}-{item.end:.2f}] {item.symbol}\n')
+                    else:
+                        text = f'{sup.speaker} {sup.text}' if sup.speaker is not None else sup.text
+                        f.write(f'[{sup.start:.2f}-{sup.end:.2f}] {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)
+                # Enhanced JSON export with word-level alignment
+                json_data = []
+                for sup in alignments:
+                    sup_dict = sup.to_dict()
+                    json_data.append(sup_dict)
+                json.dump(json_data, 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))
+                text = (
+                    f'{supervision.speaker} {supervision.text}' if supervision.speaker is not None else supervision.text
+                )
+                supervisions.append(Interval(supervision.start, supervision.end, text or ''))
+                # Extract word-level alignment using helper function
+                word_items = parse_alignment_from_supervision(supervision)
+                if word_items:
+                    for item in word_items:
+                        words.append(Interval(item.start, item.end, item.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))
+                # Add word-level timing as metadata in the subtitle text
+                word_items = parse_alignment_from_supervision(sup)
+                if word_items:
+                    for word in word_items:
+                        subs.append(
+                            pysubs2.SSAEvent(start=int(word.start * 1000), end=int(word.end * 1000), text=word.symbol)
+                        )
+                else:
+                    text = f'{sup.speaker} {sup.text}' if sup.speaker is not None else sup.text
+                    subs.append(pysubs2.SSAEvent(start=int(sup.start * 1000), end=int(sup.end * 1000), text=text or ''))
             subs.save(output_path)
 
         return output_path
+
+
+def parse_alignment_from_supervision(supervision: Any) -> Optional[List[AlignmentItem]]:
+    """
+    Extract word-level alignment items from Supervision object.
+
+    Args:
+        supervision: Supervision object with potential alignment data
+
+    Returns:
+        List of AlignmentItem objects, or None if no alignment data present
+    """
+    if not hasattr(supervision, 'alignment') or not supervision.alignment:
+        return None
+
+    if 'word' not in supervision.alignment:
+        return None
+
+    return supervision.alignment['word']

lattifai/tokenizer/__init__.py

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

lattifai/tokenizer/tokenizer.py

@@ -1,13 +1,13 @@
 import gzip
+import inspect
 import pickle
 import re
 from collections import defaultdict
-from itertools import chain
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import torch
 
-from lattifai.base_client import SyncAPIClient
+from lattifai.errors import LATTICE_DECODING_FAILURE_HELP, LatticeDecodingError
 from lattifai.io import Supervision
 from lattifai.tokenizer.phonemizer import G2Phonemizer
 
@@ -21,10 +21,13 @@ GROUPING_SEPARATOR = '✹'
 MAXIMUM_WORD_LENGTH = 40
 
 
+TokenizerT = TypeVar('TokenizerT', bound='LatticeTokenizer')
+
+
 class LatticeTokenizer:
     """Tokenizer for converting Lhotse Cut to LatticeGraph."""
 
-    def __init__(self, client_wrapper: SyncAPIClient):
+    def __init__(self, client_wrapper: Any):
         self.client_wrapper = client_wrapper
         self.words: List[str] = []
         self.g2p_model: Any = None  # Placeholder for G2P model
@@ -99,13 +102,14 @@ class LatticeTokenizer:
         # If no special pattern matches, return the original sentence
         return [sentence]
 
-    @staticmethod
+    @classmethod
     def from_pretrained(
-        client_wrapper: SyncAPIClient,
+        cls: Type[TokenizerT],
+        client_wrapper: Any,
         model_path: str,
         device: str = 'cpu',
         compressed: bool = True,
-    ):
+    ) -> TokenizerT:
         """Load tokenizer from exported binary file"""
         from pathlib import Path
 
@@ -117,7 +121,7 @@ class LatticeTokenizer:
             with open(words_model_path, 'rb') as f:
                 data = pickle.load(f)
 
-        tokenizer = LatticeTokenizer(client_wrapper=client_wrapper)
+        tokenizer = cls(client_wrapper=client_wrapper)
         tokenizer.words = data['words']
         tokenizer.dictionaries = defaultdict(list, data['dictionaries'])
         tokenizer.oov_word = data['oov_word']
@@ -179,53 +183,89 @@ class LatticeTokenizer:
         return {}
 
     def split_sentences(self, supervisions: List[Supervision], strip_whitespace=True) -> List[str]:
+        """Split supervisions into sentences using the sentence splitter.
+
+        Carefull about speaker changes.
+        """
         texts, text_len, sidx = [], 0, 0
+        speakers = []
         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)
+            if supervision.speaker:
+                speakers.append(supervision.speaker)
+                if sidx < s:
+                    text = ' '.join([sup.text for sup in supervisions[sidx:s]])
+                    texts.append(text)
+                    sidx = s
+                    text_len = len(supervision.text)
+            else:
+                if text_len >= 2000 or s == len(supervisions) - 1:
+                    if len(speakers) < len(texts) + 1:
+                        speakers.append(None)
+                    text = ' '.join([sup.text for sup in supervisions[sidx : s + 1]])
+                    texts.append(text)
+                    sidx = s + 1
+                    text_len = 0
+
+        assert len(speakers) == len(texts), f'len(speakers)={len(speakers)} != len(texts)={len(texts)}'
         sentences = self.sentence_splitter.split(texts, threshold=0.15, strip_whitespace=strip_whitespace)
 
         supervisions, remainder = [], ''
-        for _sentences in sentences:
+        for k, (_speaker, _sentences) in enumerate(zip(speakers, sentences)):
+            # Prepend remainder from previous iteration to the first sentence
+            if _sentences and remainder:
+                _sentences[0] = remainder + _sentences[0]
+                remainder = ''
+
+            if not _sentences:
+                continue
+
             # 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
+                        remainder = resplit_parts[-1] + ' '
                     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)
+                supervisions.extend(
+                    Supervision(text=text, speaker=(_speaker if s == 0 else None)) for s, text in enumerate(_sentences)
+                )
+                _speaker = None  # reset speaker after use
             else:
-                supervisions.extend(Supervision(text=s) for s in _sentences[:-1])
-                remainder += _sentences[-1] + ' '
+                supervisions.extend(
+                    Supervision(text=text, speaker=(_speaker if s == 0 else None))
+                    for s, text in enumerate(_sentences[:-1])
+                )
+                remainder = _sentences[-1] + ' ' + remainder
+                if k < len(speakers) - 1 and speakers[k + 1] is not None:  # next speaker is set
+                    supervisions.append(
+                        Supervision(text=remainder.strip(), speaker=_speaker if len(_sentences) == 1 else None)
+                    )
+                    remainder = ''
+                elif len(_sentences) == 1:
+                    if k == len(speakers) - 1:
+                        pass  # keep _speaker for the last supervision
+                    else:
+                        assert speakers[k + 1] is None
+                        speakers[k + 1] = _speaker
+                else:
+                    assert len(_sentences) > 1
+                    _speaker = None  # reset speaker if sentence not ended
 
         if remainder.strip():
-            supervisions.append(Supervision(text=remainder.strip()))
+            supervisions.append(Supervision(text=remainder.strip(), speaker=_speaker))
 
         return supervisions
 
@@ -246,14 +286,18 @@ class LatticeTokenizer:
             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))
+        return (
+            supervisions,
+            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,
+        supervisions: List[Supervision],
+        return_details: bool = False,
     ) -> List[Supervision]:
         emission, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
         response = self.client_wrapper.post(
@@ -265,20 +309,157 @@ class LatticeTokenizer:
                 'labels': labels[0],
                 'offset': offset,
                 'channel': channel,
+                'return_details': return_details,
                 'destroy_lattice': True,
             },
         )
+        if response.status_code == 422:
+            raise LatticeDecodingError(
+                lattice_id,
+                original_error=Exception(LATTICE_DECODING_FAILURE_HELP),
+            )
         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']]
+        if not result.get('success'):
+            raise Exception('Failed to detokenize the alignment results.')
+
+        alignments = [Supervision.from_dict(s) for s in result['supervisions']]
+
+        if return_details:
+            # Add emission confidence scores for segments and word-level alignments
+            _add_confidence_scores(alignments, emission, labels[0], frame_shift)
+
+        alignments = _update_alignments_speaker(supervisions, alignments)
+
+        return alignments
+
 
+class AsyncLatticeTokenizer(LatticeTokenizer):
+    async def _post_async(self, endpoint: str, **kwargs):
+        response = self.client_wrapper.post(endpoint, **kwargs)
+        if inspect.isawaitable(response):
+            return await response
+        return response
 
-# 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)
+    async 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 = await self._post_async(
+            '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 (
+            supervisions,
+            lattice_id,
+            (result['lattice_graph'], result['final_state'], result.get('acoustic_scale', 1.0)),
+        )
+
+    async def detokenize(
+        self,
+        lattice_id: str,
+        lattice_results: Tuple[torch.Tensor, Any, Any, float, float],
+        supervisions: List[Supervision],
+        return_details: bool = False,
+    ) -> List[Supervision]:
+        emission, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
+        response = await self._post_async(
+            '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,
+                'return_details': return_details,
+                'destroy_lattice': True,
+            },
+        )
+        if response.status_code == 422:
+            raise LatticeDecodingError(
+                lattice_id,
+                original_error=Exception(LATTICE_DECODING_FAILURE_HELP),
+            )
+        if response.status_code != 200:
+            raise Exception(f'Failed to detokenize lattice: {response.text}')
+
+        result = response.json()
+        if not result.get('success'):
+            return Exception('Failed to detokenize the alignment results.')
+
+        alignments = [Supervision.from_dict(s) for s in result['supervisions']]
+
+        if return_details:
+            # Add emission confidence scores for segments and word-level alignments
+            _add_confidence_scores(alignments, emission, labels[0], frame_shift)
+
+        alignments = _update_alignments_speaker(supervisions, alignments)
+
+        return alignments
+
+
+def _add_confidence_scores(
+    supervisions: List[Supervision],
+    emission: torch.Tensor,
+    labels: List[int],
+    frame_shift: float,
+) -> None:
+    """
+    Add confidence scores to supervisions and their word-level alignments.
+
+    This function modifies supervisions in-place by:
+    1. Computing segment-level confidence scores based on emission probabilities
+    2. Computing word-level confidence scores for each aligned word
+
+    Args:
+        supervisions: List of Supervision objects to add scores to (modified in-place)
+        emission: Emission tensor with shape [batch, time, vocab_size]
+        labels: Token labels corresponding to aligned tokens
+        frame_shift: Frame shift in seconds for converting frames to time
+    """
+    tokens = torch.tensor(labels, dtype=torch.int64, device=emission.device)
+
+    for supervision in supervisions:
+        start_frame = int(supervision.start / frame_shift)
+        end_frame = int(supervision.end / frame_shift)
+
+        # Compute segment-level confidence
+        probabilities = emission[0, start_frame:end_frame].softmax(dim=-1)
+        aligned = probabilities[range(0, end_frame - start_frame), tokens[start_frame:end_frame]]
+        diffprobs = (probabilities.max(dim=-1).values - aligned).cpu()
+        supervision.score = round(1.0 - diffprobs.mean().item(), ndigits=4)
+
+        # Compute word-level confidence if alignment exists
+        if hasattr(supervision, 'alignment') and supervision.alignment:
+            words = supervision.alignment.get('word', [])
+            for w, item in enumerate(words):
+                start = int(item.start / frame_shift) - start_frame
+                end = int(item.end / frame_shift) - start_frame
+                words[w] = item._replace(score=round(1.0 - diffprobs[start:end].mean().item(), ndigits=4))
+
+
+def _update_alignments_speaker(supervisions: List[Supervision], alignments: List[Supervision]) -> List[Supervision]:
+    """
+    Update the speaker attribute for a list of supervisions.
+
+    Args:
+        supervisions: List of Supervision objects to get speaker info from
+        alignments: List of aligned Supervision objects to update speaker info to
+    """
+    for supervision, alignment in zip(supervisions, alignments):
+        alignment.speaker = supervision.speaker
+    return alignments

lattifai/utils.py

@@ -0,0 +1,133 @@
+"""Shared utility helpers for the LattifAI SDK."""
+
+import os
+from datetime import datetime, timedelta
+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:
+    """Get the path for the cache marker file with current date."""
+    today = datetime.now().strftime('%Y%m%d')
+    return cache_dir / f'.done{today}'
+
+
+def _is_cache_valid(cache_dir: Path) -> bool:
+    """Check if cached model is valid (exists and not older than 1 days)."""
+    if not cache_dir.exists():
+        return False
+
+    # Find any .done* marker files
+    marker_files = list(cache_dir.glob('.done*'))
+    if not marker_files:
+        return False
+
+    # Get the most recent marker file
+    latest_marker = max(marker_files, key=lambda p: p.stat().st_mtime)
+
+    # Extract date from marker filename (format: .doneYYYYMMDD)
+    try:
+        date_str = latest_marker.name.replace('.done', '')
+        marker_date = datetime.strptime(date_str, '%Y%m%d')
+        # Check if marker is older than 1 days
+        if datetime.now() - marker_date > timedelta(days=1):
+            return False
+        return True
+    except (ValueError, IndexError):
+        # Invalid marker file format, treat as invalid cache
+        return False
+
+
+def _create_cache_marker(cache_dir: Path) -> None:
+    """Create a cache marker file with current date and clean old markers."""
+    # Remove old marker files
+    for old_marker in cache_dir.glob('.done*'):
+        old_marker.unlink(missing_ok=True)
+
+    # Create new marker file
+    marker_path = _get_cache_marker_path(cache_dir)
+    marker_path.touch()
+
+
+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
+
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.constants import HF_HUB_CACHE
+    from huggingface_hub.errors import LocalEntryNotFoundError
+
+    # Determine cache directory for this model
+    cache_dir = Path(HF_HUB_CACHE) / f'models--{model_name_or_path.replace("/", "--")}'
+
+    # Check if we have a valid cached version
+    if _is_cache_valid(cache_dir):
+        # Return the snapshot path (latest version)
+        snapshots_dir = cache_dir / 'snapshots'
+        if snapshots_dir.exists():
+            snapshot_dirs = [d for d in snapshots_dir.iterdir() if d.is_dir()]
+            if snapshot_dirs:
+                # Return the most recent snapshot
+                latest_snapshot = max(snapshot_dirs, key=lambda p: p.stat().st_mtime)
+                return str(latest_snapshot)
+
+    try:
+        downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
+        _create_cache_marker(cache_dir)
+        return downloaded_path
+    except LocalEntryNotFoundError:
+        try:
+            os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+            downloaded_path = snapshot_download(repo_id=model_name_or_path, repo_type='model')
+            _create_cache_marker(cache_dir)
+            return downloaded_path
+        except Exception as e:  # pragma: no cover - bubble up for caller context
+            raise ModelLoadError(model_name_or_path, original_error=e)
+    except Exception as e:  # pragma: no cover - unexpected download issue
+        raise ModelLoadError(model_name_or_path, original_error=e)
+
+
+def _select_device(device: Optional[str]) -> str:
+    """Select best available torch device when not explicitly provided."""
+    if device:
+        return device
+
+    import torch
+
+    detected = 'cpu'
+    if torch.backends.mps.is_available():
+        detected = 'mps'
+    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)