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

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_mode:
-                    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/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/sentence_splitter.py

@@ -1,16 +1,15 @@
 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
 
-END_PUNCTUATION = '.!?"]。！？"】'
-
 
 class SentenceSplitter:
     """Lazy-initialized sentence splitter using wtpsplit."""
 
-    def __init__(self, device: str = "cpu", model_hub: Optional[str] = None, lazy_init: bool = True):
+    def __init__(self, device: str = "cpu", model_hub: Optional[str] = "modelscope", lazy_init: bool = True):
         """Initialize sentence splitter with lazy loading.
 
         Args:
@@ -19,9 +18,8 @@ class SentenceSplitter:
         """
         self.device = device
         self.model_hub = model_hub
-        if lazy_init:
-            self._splitter = None
-        else:
+        self._splitter = None
+        if not lazy_init:
             self._init_splitter()
 
     def _init_splitter(self):
@@ -56,6 +54,121 @@ class SentenceSplitter:
             )
         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]:
         """
@@ -150,15 +263,22 @@ class SentenceSplitter:
             elif text_len >= 2000 or is_last:
                 flush_segment(s, None)
 
-        assert len(speakers) == len(texts), f"len(speakers)={len(speakers)} != len(texts)={len(texts)}"
+        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)
 
-        supervisions, remainder = [], ""
+        # 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
@@ -188,32 +308,43 @@ class SentenceSplitter:
                     continue
 
             if any(_sentences[-1].endswith(ep) for ep in END_PUNCTUATION):
-                supervisions.extend(
-                    Supervision(text=text, speaker=(_speaker if s == 0 else None)) for s, text in enumerate(_sentences)
+                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:
-                supervisions.extend(
-                    Supervision(text=text, speaker=(_speaker if s == 0 else None))
-                    for s, text in enumerate(_sentences[:-1])
+                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
-                    supervisions.append(
-                        Supervision(text=remainder.strip(), speaker=_speaker if len(_sentences) == 1 else None)
-                    )
+                    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:
-                        assert speakers[k + 1] is None
                         speakers[k + 1] = _speaker
-                else:
-                    assert len(_sentences) > 1
+                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():
-            supervisions.append(Supervision(text=remainder.strip(), speaker=_speaker))
+            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 supervisions
+        return result_supervisions