lattifai 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

lattifai/__init__.py

@@ -1,5 +1,4 @@
 import os
-import sys
 import warnings
 from importlib.metadata import version
 
@@ -52,29 +51,6 @@ except Exception:
     __version__ = "0.1.0"  # fallback version
 
 
-# Check and auto-install k2py if not present
-def _check_and_install_k2py():
-    """Check if k2py is installed and attempt to install it if not."""
-    try:
-        import k2py
-    except ImportError:
-        import subprocess
-
-        print("k2py is not installed. Attempting to install k2py...")
-        try:
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "k2py"])
-            import k2py  # Try importing again after installation
-
-            print("k2py installed successfully.")
-        except Exception as e:
-            warnings.warn(f"Failed to install k2py automatically. Please install it manually. Error: {e}")
-    return True
-
-
-# Auto-install k2py on first import
-_check_and_install_k2py()
-
-
 __all__ = [
     # Client classes
     "LattifAI",

lattifai/alignment/__init__.py

@@ -2,5 +2,14 @@
 
 from .lattice1_aligner import Lattice1Aligner
 from .segmenter import Segmenter
+from .sentence_splitter import SentenceSplitter
+from .text_align import align_supervisions_and_transcription
+from .tokenizer import tokenize_multilingual_text
 
-__all__ = ["Lattice1Aligner", "Segmenter"]
+__all__ = [
+    "Lattice1Aligner",
+    "Segmenter",
+    "SentenceSplitter",
+    "align_supervisions_and_transcription",
+    "tokenize_multilingual_text",
+]

lattifai/alignment/lattice1_aligner.py

@@ -1,6 +1,6 @@
 """Lattice-1 Aligner implementation."""
 
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
 import colorful
 import numpy as np
@@ -16,11 +16,22 @@ from lattifai.errors import (
 from lattifai.utils import _resolve_model_path, safe_print
 
 from .lattice1_worker import _load_worker
+from .text_align import TextAlignResult
 from .tokenizer import _load_tokenizer
 
 ClientType = Any
 
 
+def _extract_text_for_error(supervisions: Union[list, tuple]) -> str:
+    """Extract text from supervisions for error messages."""
+    if not supervisions:
+        return ""
+    # TextAlignResult is a tuple: (caption_sups, transcript_sups, ...)
+    if isinstance(supervisions, tuple):
+        supervisions = supervisions[0] or supervisions[1] or []
+    return " ".join(s.text for s in supervisions if s and s.text)
+
+
 class Lattice1Aligner(object):
     """Synchronous LattifAI client with config-driven architecture."""
 
@@ -79,7 +90,7 @@ class Lattice1Aligner(object):
     def alignment(
         self,
         audio: AudioData,
-        supervisions: List[Supervision],
+        supervisions: Union[List[Supervision], TextAlignResult],
         split_sentence: Optional[bool] = False,
         return_details: Optional[bool] = False,
         emission: Optional[np.ndarray] = None,
@@ -102,69 +113,66 @@ class Lattice1Aligner(object):
             AlignmentError: If audio alignment fails
             LatticeDecodingError: If lattice decoding fails
         """
+        # Step 2: Create lattice graph
+        if verbose:
+            safe_print(colorful.cyan("🔗 Step 2: Creating lattice graph from segments"))
         try:
+            supervisions, lattice_id, lattice_graph = self.tokenizer.tokenize(
+                supervisions, split_sentence=split_sentence, boost=self.config.boost
+            )
             if verbose:
-                safe_print(colorful.cyan("🔗 Step 2: Creating lattice graph from segments"))
-            try:
-                supervisions, lattice_id, lattice_graph = self.tokenizer.tokenize(
-                    supervisions, split_sentence=split_sentence
-                )
-                if verbose:
-                    safe_print(colorful.green(f"         ✓ Generated lattice graph with ID: {lattice_id}"))
-            except Exception as e:
-                text_content = " ".join([sup.text for sup in supervisions]) if supervisions else ""
-                raise LatticeEncodingError(text_content, original_error=e)
-
-            if verbose:
-                safe_print(colorful.cyan(f"🔍 Step 3: Searching lattice graph with media: {audio}"))
-                if audio.streaming_chunk_secs:
-                    safe_print(
-                        colorful.yellow(
-                            f"         ⚡Using streaming mode with {audio.streaming_chunk_secs}s (chunk duration)"
-                        )
+                safe_print(colorful.green(f"         ✓ Generated lattice graph with ID: {lattice_id}"))
+        except Exception as e:
+            text_content = _extract_text_for_error(supervisions)
+            raise LatticeEncodingError(text_content, original_error=e)
+
+        # Step 3: Search lattice graph
+        if verbose:
+            safe_print(colorful.cyan(f"🔍 Step 3: Searching lattice graph with media: {audio}"))
+            if audio.streaming_mode:
+                safe_print(
+                    colorful.yellow(
+                        f"         ⚡Using streaming mode with {audio.streaming_chunk_secs}s (chunk duration)"
                     )
-            try:
-                lattice_results = self.worker.alignment(
-                    audio,
-                    lattice_graph,
-                    emission=emission,
-                    offset=offset,
-                )
-                if verbose:
-                    safe_print(colorful.green("         ✓ Lattice search completed"))
-            except Exception as e:
-                raise AlignmentError(
-                    f"Audio alignment failed for {audio}",
-                    media_path=str(audio),
-                    context={"original_error": str(e)},
                 )
-
+        try:
+            lattice_results = self.worker.alignment(
+                audio,
+                lattice_graph,
+                emission=emission,
+                offset=offset,
+            )
             if verbose:
-                safe_print(colorful.cyan("🎯 Step 4: Decoding lattice results to aligned segments"))
-            try:
-                alignments = self.tokenizer.detokenize(
-                    lattice_id,
-                    lattice_results,
-                    supervisions=supervisions,
-                    return_details=return_details,
-                    start_margin=self.config.start_margin,
-                    end_margin=self.config.end_margin,
-                )
-                if verbose:
-                    safe_print(colorful.green(f"         ✓ Successfully aligned {len(alignments)} segments"))
-            except LatticeDecodingError as e:
-                safe_print(colorful.red("         x Failed to decode lattice alignment results"))
-                raise e
-            except Exception as e:
-                safe_print(colorful.red("         x Failed to decode lattice alignment results"))
-                raise LatticeDecodingError(lattice_id, original_error=e)
-
-            return (supervisions, alignments)
-
-        except (LatticeEncodingError, AlignmentError, LatticeDecodingError):
+                safe_print(colorful.green("         ✓ Lattice search completed"))
+        except Exception as e:
+            raise AlignmentError(
+                f"Audio alignment failed for {audio}",
+                media_path=str(audio),
+                context={"original_error": str(e)},
+            )
+
+        # Step 4: Decode lattice results
+        if verbose:
+            safe_print(colorful.cyan("🎯 Step 4: Decoding lattice results to aligned segments"))
+        try:
+            alignments = self.tokenizer.detokenize(
+                lattice_id,
+                lattice_results,
+                supervisions=supervisions,
+                return_details=return_details,
+                start_margin=self.config.start_margin,
+                end_margin=self.config.end_margin,
+            )
+            if verbose:
+                safe_print(colorful.green(f"         ✓ Successfully aligned {len(alignments)} segments"))
+        except LatticeDecodingError:
+            safe_print(colorful.red("         x Failed to decode lattice alignment results"))
             raise
         except Exception as e:
-            raise e
+            safe_print(colorful.red("         x Failed to decode lattice alignment results"))
+            raise LatticeDecodingError(lattice_id, original_error=e)
+
+        return (supervisions, alignments)
 
     def profile(self) -> None:
         """Print profiling statistics."""

lattifai/alignment/lattice1_worker.py

@@ -7,8 +7,6 @@ from typing import Any, Dict, Optional, Tuple
 import colorful
 import numpy as np
 import onnxruntime as ort
-from lhotse import FbankConfig
-from lhotse.features.kaldi.layers import Wav2LogFilterBank
 from lhotse.utils import Pathlike
 from tqdm import tqdm
 
@@ -159,10 +157,7 @@ class Lattice1Worker:
             DependencyError: If required dependencies are missing
             AlignmentError: If alignment process fails
         """
-        try:
-            import k2py as k2
-        except ImportError:
-            raise DependencyError("k2py", install_command="pip install k2py")
+        import k2py as k2
 
         lattice_graph_str, final_state, acoustic_scale = lattice_graph
 

lattifai/alignment/punctuation.py

@@ -0,0 +1,38 @@
+# Multilingual punctuation characters (no duplicates)
+PUNCTUATION = (
+    # ASCII punctuation
+    "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
+    # Chinese/CJK punctuation
+    "。，、？！：；·￥…—～“”‘’"
+    "《》〈〉【】〔〕〖〗（）"  # Brackets
+    # Japanese punctuation (unique only)
+    "「」『』・"
+    # CJK punctuation (unique only)
+    "〃〆〇〒〓〘〙〚〛〜〝〞〟"
+    # Arabic punctuation
+    "،؛؟"
+    # Thai punctuation
+    "๏๚๛"
+    # Hebrew punctuation
+    "־׀׃׆"
+    # Other common punctuation
+    "¡¿"  # Spanish inverted marks
+    "«»‹›"  # Guillemets
+    "‐‑‒–―"  # Dashes (excluding — already above)
+    "‚„"  # Low quotation marks
+    "†‡•‣"  # Daggers and bullets
+    "′″‴"  # Prime marks
+    "‰‱"  # Per mille
+)
+PUNCTUATION_SPACE = PUNCTUATION + " "
+STAR_TOKEN = "※"
+
+# End of sentence punctuation marks (multilingual)
+# - ASCII: .!?"']）
+# - Chinese/CJK: 。！？"】」』〗〙〛 (including right double quote U+201D)
+# - Japanese: ｡ (halfwidth period)
+# - Arabic: ؟
+# - Ellipsis: …
+END_PUNCTUATION = ".!?\"']）。！？\u201d】」』〗〙〛｡؟…"
+
+GROUPING_SEPARATOR = "✹"

lattifai/alignment/segmenter.py

@@ -9,7 +9,7 @@ from lattifai.caption import Caption, Supervision
 from lattifai.config import AlignmentConfig
 from lattifai.utils import safe_print
 
-from .tokenizer import END_PUNCTUATION
+from .sentence_splitter import END_PUNCTUATION
 
 
 class Segmenter:

lattifai/alignment/sentence_splitter.py

@@ -0,0 +1,350 @@
+import re
+from typing import List, Optional
+
+from lattifai.alignment.punctuation import END_PUNCTUATION
+from lattifai.caption import Supervision
+from lattifai.utils import _resolve_model_path
+
+
+class SentenceSplitter:
+    """Lazy-initialized sentence splitter using wtpsplit."""
+
+    def __init__(self, device: str = "cpu", model_hub: Optional[str] = "modelscope", lazy_init: bool = True):
+        """Initialize sentence splitter with lazy loading.
+
+        Args:
+            device: Device to run the model on (cpu, cuda, mps)
+            model_hub: Model hub to use (None for huggingface, "modelscope" for modelscope)
+        """
+        self.device = device
+        self.model_hub = model_hub
+        self._splitter = None
+        if not lazy_init:
+            self._init_splitter()
+
+    def _init_splitter(self):
+        """Initialize the sentence splitter model on first use."""
+        if self._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")
+
+        if self.model_hub == "modelscope":
+            downloaded_path = _resolve_model_path("LattifAI/OmniTokenizer", model_hub="modelscope")
+            sat = SaT(
+                f"{downloaded_path}/sat-3l-sm",
+                tokenizer_name_or_path=f"{downloaded_path}/xlm-roberta-base",
+                ort_providers=providers + ["CPUExecutionProvider"],
+            )
+        else:
+            sat_path = _resolve_model_path("segment-any-text/sat-3l-sm", model_hub="huggingface")
+            sat = SaT(
+                sat_path,
+                tokenizer_name_or_path="facebookAI/xlm-roberta-base",
+                hub_prefix="segment-any-text",
+                ort_providers=providers + ["CPUExecutionProvider"],
+            )
+        self._splitter = sat
+
+    @staticmethod
+    def _distribute_time_info(
+        input_supervisions: List[Supervision],
+        split_texts: List[str],
+    ) -> List[Supervision]:
+        """Distribute time information from input supervisions to split sentences.
+
+        Args:
+            input_supervisions: Original supervisions with time information
+            split_texts: List of split sentence texts
+
+        Returns:
+            List of Supervision objects with distributed time information.
+            Custom attributes are inherited from first_sup with conflict markers.
+        """
+        if not input_supervisions:
+            return [Supervision(text=text, id="", recording_id="", start=0, duration=0) for text in split_texts]
+
+        # Build concatenated input text
+        input_text = " ".join(sup.text for sup in input_supervisions)
+
+        # Pre-compute supervision position mapping for O(1) lookup
+        # Format: [(start_pos, end_pos, supervision), ...]
+        sup_ranges = []
+        char_pos = 0
+        for sup in input_supervisions:
+            sup_start = char_pos
+            sup_end = char_pos + len(sup.text)
+            sup_ranges.append((sup_start, sup_end, sup))
+            char_pos = sup_end + 1  # +1 for space separator
+
+        # Process each split text
+        result = []
+        search_start = 0
+        sup_idx = 0  # Track current supervision index to skip processed ones
+
+        for split_text in split_texts:
+            text_start = input_text.find(split_text, search_start)
+            if text_start == -1:
+                raise ValueError(f"Could not find split text '{split_text}' in input supervisions.")
+
+            text_end = text_start + len(split_text)
+            search_start = text_end
+
+            # Find overlapping supervisions, starting from last used index
+            first_sup = None
+            last_sup = None
+            first_char_idx = None
+            last_char_idx = None
+            overlapping_customs = []  # Track all custom dicts for conflict detection
+
+            # Start from sup_idx, which is the first supervision that might overlap
+            for i in range(sup_idx, len(sup_ranges)):
+                sup_start, sup_end, sup = sup_ranges[i]
+
+                # Skip if no overlap (before text_start)
+                if sup_end <= text_start:
+                    sup_idx = i + 1  # Update starting point for next iteration
+                    continue
+
+                # Stop if no overlap (after text_end)
+                if sup_start >= text_end:
+                    break
+
+                # Found overlap
+                if first_sup is None:
+                    first_sup = sup
+                    first_char_idx = max(0, text_start - sup_start)
+
+                last_sup = sup
+                last_char_idx = min(len(sup.text) - 1, text_end - 1 - sup_start)
+
+                # Collect custom dict for conflict detection
+                if getattr(sup, "custom", None):
+                    overlapping_customs.append(sup.custom)
+
+            if first_sup is None or last_sup is None:
+                raise ValueError(f"Could not find supervisions for split text: {split_text}")
+
+            # Calculate timing
+            start_time = first_sup.start + (first_char_idx / len(first_sup.text)) * first_sup.duration
+            end_time = last_sup.start + ((last_char_idx + 1) / len(last_sup.text)) * last_sup.duration
+
+            # Inherit custom from first_sup, mark conflicts if multiple sources
+            merged_custom = None
+            if overlapping_customs:
+                # Start with first_sup's custom (inherit strategy)
+                merged_custom = overlapping_customs[0].copy() if overlapping_customs[0] else {}
+
+                # Detect conflicts if multiple overlapping supervisions have different custom values
+                if len(overlapping_customs) > 1:
+                    has_conflict = False
+                    for other_custom in overlapping_customs[1:]:
+                        if other_custom and other_custom != overlapping_customs[0]:
+                            has_conflict = True
+                            break
+
+                    if has_conflict:
+                        # Mark that this supervision spans multiple sources with different customs
+                        merged_custom["_split_from_multiple"] = True
+                        merged_custom["_source_count"] = len(overlapping_customs)
+
+            result.append(
+                Supervision(
+                    id="",
+                    text=split_text,
+                    start=start_time,
+                    duration=end_time - start_time,
+                    recording_id=first_sup.recording_id,
+                    custom=merged_custom,
+                )
+            )
+
+        return result
+
+    @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]
+
+    def split_sentences(self, supervisions: List[Supervision], strip_whitespace=True) -> List[Supervision]:
+        """Split supervisions into sentences using the sentence splitter.
+
+        Careful about speaker changes.
+
+        Args:
+            supervisions: List of Supervision objects to split
+            strip_whitespace: Whether to strip whitespace from split sentences
+
+        Returns:
+            List of Supervision objects with split sentences
+        """
+        self._init_splitter()
+
+        texts, speakers = [], []
+        text_len, sidx = 0, 0
+
+        def flush_segment(end_idx: int, speaker: Optional[str] = None):
+            """Flush accumulated text from sidx to end_idx with given speaker."""
+            nonlocal text_len, sidx
+            if sidx <= end_idx:
+                if len(speakers) < len(texts) + 1:
+                    speakers.append(speaker)
+                text = " ".join(sup.text for sup in supervisions[sidx : end_idx + 1])
+                texts.append(text)
+                sidx = end_idx + 1
+                text_len = 0
+
+        for s, supervision in enumerate(supervisions):
+            text_len += len(supervision.text)
+            is_last = s == len(supervisions) - 1
+
+            if supervision.speaker:
+                # Flush previous segment without speaker (if any)
+                if sidx < s:
+                    flush_segment(s - 1, None)
+                    text_len = len(supervision.text)
+
+                # Check if we should flush this speaker's segment now
+                next_has_speaker = not is_last and supervisions[s + 1].speaker
+                if is_last or next_has_speaker:
+                    flush_segment(s, supervision.speaker)
+                else:
+                    speakers.append(supervision.speaker)
+
+            elif text_len >= 2000 or is_last:
+                flush_segment(s, None)
+
+        if len(speakers) != len(texts):
+            raise ValueError(f"len(speakers)={len(speakers)} != len(texts)={len(texts)}")
+        sentences = self._splitter.split(texts, threshold=0.15, strip_whitespace=strip_whitespace, batch_size=8)
+
+        # First pass: collect all split texts with their speakers
+        split_texts_with_speakers = []
+        remainder = ""
+        remainder_speaker = None
+
+        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]
+                _speaker = remainder_speaker if remainder_speaker else _speaker
+                remainder = ""
+                remainder_speaker = None
+
+            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] + " "
+                    processed_sentences.extend(resplit_parts[:-1])
+                else:
+                    processed_sentences.extend(resplit_parts)
+            _sentences = processed_sentences
+
+            if not _sentences:
+                if remainder:
+                    _sentences, remainder = [remainder.strip()], ""
+                else:
+                    continue
+
+            if any(_sentences[-1].endswith(ep) for ep in END_PUNCTUATION):
+                split_texts_with_speakers.extend(
+                    (text, _speaker if s == 0 else None) for s, text in enumerate(_sentences)
+                )
+                _speaker = None  # reset speaker after use
+            else:
+                split_texts_with_speakers.extend(
+                    (text, _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
+                    split_texts_with_speakers.append((remainder.strip(), _speaker if len(_sentences) == 1 else None))
+                    remainder = ""
+                    remainder_speaker = None
+                elif len(_sentences) == 1:
+                    remainder_speaker = _speaker
+                    if k == len(speakers) - 1:
+                        pass  # keep _speaker for the last supervision
+                    elif speakers[k + 1] is not None:
+                        raise ValueError(f"Expected speakers[{k + 1}] to be None, got {speakers[k + 1]}")
+                    else:
+                        speakers[k + 1] = _speaker
+                elif len(_sentences) > 1:
+                    _speaker = None  # reset speaker if sentence not ended
+                    remainder_speaker = None
+                else:
+                    raise ValueError(f"Unexpected state: len(_sentences)={len(_sentences)}")
+
+        if remainder.strip():
+            split_texts_with_speakers.append((remainder.strip(), remainder_speaker))
+
+        # Second pass: distribute time information
+        split_texts = [text for text, _ in split_texts_with_speakers]
+        result_supervisions = self._distribute_time_info(supervisions, split_texts)
+
+        # Third pass: add speaker information
+        for sup, (_, speaker) in zip(result_supervisions, split_texts_with_speakers):
+            if speaker:
+                sup.speaker = speaker
+
+        return result_supervisions