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

lattifai/_init.py

@@ -0,0 +1,20 @@
+"""Environment configuration for LattifAI.
+
+Import this module early to suppress warnings before other imports.
+
+Usage:
+    import lattifai._init  # noqa: F401
+    from lattifai.client import LattifAI
+"""
+
+import os
+import warnings
+
+# Suppress SWIG deprecation warnings before any imports
+warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*")
+
+# Suppress PyTorch transformer nested tensor warning
+warnings.filterwarnings("ignore", category=UserWarning, message=".*enable_nested_tensor.*")
+
+# Disable tokenizers parallelism warning
+os.environ["TOKENIZERS_PARALLELISM"] = "false"

lattifai/alignment/__init__.py

@@ -1,15 +1,14 @@
 """Alignment module for LattifAI forced alignment."""
 
+from lattifai.caption import SentenceSplitter
+
 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",
     "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, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import colorful
 import numpy as np
@@ -118,7 +118,10 @@ class Lattice1Aligner(object):
             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
+                supervisions,
+                split_sentence=split_sentence,
+                boost=self.config.boost,
+                transition_penalty=self.config.transition_penalty,
             )
             if verbose:
                 safe_print(colorful.green(f"         ✓ Generated lattice graph with ID: {lattice_id}"))
@@ -162,12 +165,34 @@ class Lattice1Aligner(object):
                 return_details=return_details,
                 start_margin=self.config.start_margin,
                 end_margin=self.config.end_margin,
+                check_sanity=True,
             )
             if verbose:
                 safe_print(colorful.green(f"         ✓ Successfully aligned {len(alignments)} segments"))
-        except LatticeDecodingError:
+        except LatticeDecodingError as e:
             safe_print(colorful.red("         x Failed to decode lattice alignment results"))
-            raise
+            _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,
+                check_sanity=False,
+            )
+            # Check for score anomalies (media-text mismatch)
+            anomaly = _detect_score_anomalies(_alignments)
+            if anomaly:
+                anomaly_str = _format_anomaly_warning(anomaly)
+                del _alignments
+                raise LatticeDecodingError(
+                    lattice_id,
+                    message=colorful.yellow("Score anomaly detected - media and text mismatch:\n" + anomaly_str),
+                    skip_help=True,  # anomaly info is more specific than default help
+                )
+            else:
+                del _alignments
+                raise e
         except Exception as e:
             safe_print(colorful.red("         x Failed to decode lattice alignment results"))
             raise LatticeDecodingError(lattice_id, original_error=e)
@@ -177,3 +202,91 @@ class Lattice1Aligner(object):
     def profile(self) -> None:
         """Print profiling statistics."""
         self.worker.profile()
+
+
+def _detect_score_anomalies(
+    alignments: List[Supervision],
+    drop_threshold: float = 0.08,
+    window_size: int = 5,
+) -> Optional[Dict[str, Any]]:
+    """Detect score anomalies indicating alignment mismatch.
+
+    Compares average of window_size segments before vs after each position.
+    When the drop is significant, it indicates the audio doesn't match
+    the text starting at that position.
+
+    Args:
+        alignments: List of aligned supervisions with scores
+        drop_threshold: Minimum drop between before/after averages to trigger
+        window_size: Number of segments to average on each side
+
+    Returns:
+        Dict with anomaly info if found, None otherwise
+    """
+    scores = [s.score for s in alignments if s.score is not None]
+    if len(scores) < window_size * 2:
+        return None
+
+    for i in range(window_size, len(scores) - window_size):
+        before_avg = np.mean(scores[i - window_size : i])
+        after_avg = np.mean(scores[i : i + window_size])
+        drop = before_avg - after_avg
+
+        # Trigger: significant drop between before and after windows
+        if drop > drop_threshold:
+            # Find the exact mutation point (largest single-step drop)
+            max_drop = 0
+            mutation_idx = i
+            for j in range(i - 1, min(i + window_size, len(scores) - 1)):
+                single_drop = scores[j] - scores[j + 1]
+                if single_drop > max_drop:
+                    max_drop = single_drop
+                    mutation_idx = j + 1
+
+            # Segments: last normal + anomaly segments
+            last_normal = alignments[mutation_idx - 1] if mutation_idx > 0 else None
+            anomaly_segments = [
+                alignments[j] for j in range(mutation_idx, min(mutation_idx + window_size, len(alignments)))
+            ]
+
+            return {
+                "mutation_index": mutation_idx,
+                "before_avg": round(before_avg, 4),
+                "after_avg": round(after_avg, 4),
+                "window_drop": round(drop, 4),
+                "mutation_drop": round(max_drop, 4),
+                "last_normal": last_normal,
+                "segments": anomaly_segments,
+            }
+
+    return None
+
+
+def _format_anomaly_warning(anomaly: Dict[str, Any]) -> str:
+    """Format anomaly detection result as warning message."""
+    lines = [
+        f"⚠️  Score anomaly detected at segment #{anomaly['mutation_index']}",
+        f"    Window avg: {anomaly['before_avg']:.4f} → {anomaly['after_avg']:.4f} (drop: {anomaly['window_drop']:.4f})",  # noqa: E501
+        f"    Mutation drop: {anomaly['mutation_drop']:.4f}",
+        "",
+    ]
+
+    # Show last normal segment
+    if anomaly.get("last_normal"):
+        seg = anomaly["last_normal"]
+        text_preview = seg.text[:50] + "..." if len(seg.text) > 50 else seg.text
+        lines.append(f'    [{seg.start:.2f}s-{seg.end:.2f}s] score={seg.score:.4f} "{text_preview}"')
+
+    # Separator - mutation point
+    lines.append("    " + "─" * 60)
+    lines.append(f"    ⬇️  MUTATION: The following {len(anomaly['segments'])}+ segments don't match audio")
+    lines.append("    " + "─" * 60)
+
+    # Show anomaly segments
+    for seg in anomaly["segments"]:
+        text_preview = seg.text[:50] + "..." if len(seg.text) > 50 else seg.text
+        lines.append(f'    [{seg.start:.2f}s-{seg.end:.2f}s] score={seg.score:.4f} "{text_preview}"')
+
+    lines.append("")
+    lines.append("    Possible causes: Transcription error, missing content, or wrong audio region")
+    return "\n".join(lines)

lattifai/alignment/lattice1_worker.py

@@ -35,6 +35,8 @@ class Lattice1Worker:
         sess_options.intra_op_num_threads = num_threads  # CPU cores
         sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
         sess_options.add_session_config_entry("session.intra_op.allow_spinning", "0")
+        # Suppress CoreMLExecutionProvider warnings about partial graph support
+        sess_options.log_severity_level = 3  # ERROR level only
 
         acoustic_model_path = f"{model_path}/acoustic_opt.onnx"
 
@@ -191,12 +193,17 @@ class Lattice1Worker:
                 float(output_beam),
                 int(min_active_states),
                 int(max_active_states),
+                allow_partial=True,
             )
 
-            # Streaming mode
+            # Streaming mode with confidence score accumulation
             total_duration = audio.duration
             total_minutes = int(total_duration / 60.0)
 
+            max_probs = []
+            aligned_probs = []
+            prev_labels_len = 0
+
             with tqdm(
                 total=total_minutes,
                 desc=f"Processing audio ({total_minutes} min)",
@@ -208,13 +215,43 @@ class Lattice1Worker:
                     chunk_emission = self.emission(chunk.ndarray, acoustic_scale=acoustic_scale)
                     intersecter.decode(chunk_emission[0])
 
+                    __start = time.time()
+                    # Get partial labels and compute confidence stats for this chunk
+                    partial_labels = intersecter.get_partial_labels()
+                    chunk_len = chunk_emission.shape[1]
+
+                    # Get labels for current chunk (new labels since last chunk)
+                    chunk_labels = partial_labels[prev_labels_len : prev_labels_len + chunk_len]
+                    prev_labels_len = len(partial_labels)
+
+                    # Compute emission-based confidence stats
+                    probs = np.exp(chunk_emission[0])  # [T, V]
+                    max_probs.append(np.max(probs, axis=-1))  # [T]
+
+                    # Handle case where chunk_labels length might differ from chunk_len
+                    if len(chunk_labels) == chunk_len:
+                        aligned_probs.append(probs[np.arange(chunk_len), chunk_labels])
+                    else:
+                        # Fallback: use max probs as aligned probs (approximate)
+                        aligned_probs.append(np.max(probs, axis=-1))
+
+                    del chunk_emission, probs  # Free memory
+                    self.timings["align_>labels"] += time.time() - __start
+
                     # Update progress
                     chunk_duration = int(chunk.duration / 60.0)
                     pbar.update(chunk_duration)
 
-            emission_result = None
+            # Build emission_stats for confidence calculation
+            emission_stats = {
+                "max_probs": np.concatenate(max_probs),
+                "aligned_probs": np.concatenate(aligned_probs),
+            }
+
             # Get results from intersecter
+            __start = time.time()
             results, labels = intersecter.finish()
+            self.timings["align_>finish"] += time.time() - __start
         else:
             # Batch mode
             if emission is None:
@@ -230,13 +267,19 @@ class Lattice1Worker:
                 float(output_beam),
                 int(min_active_states),
                 int(max_active_states),
+                allow_partial=True,
             )
-            emission_result = emission
+            # Compute emission_stats from full emission (same format as streaming)
+            probs = np.exp(emission[0])  # [T, V]
+            emission_stats = {
+                "max_probs": np.max(probs, axis=-1),  # [T]
+                "aligned_probs": probs[np.arange(probs.shape[0]), labels[0]],  # [T]
+            }
 
         self.timings["align_segments"] += time.time() - _start
 
         channel = 0
-        return emission_result, results, labels, self.frame_shift, offset, channel  # frame_shift=20ms
+        return emission_stats, results, labels, self.frame_shift, offset, channel  # frame_shift=20ms
 
     def profile(self) -> None:
         """Print formatted profiling statistics."""

lattifai/alignment/segmenter.py

@@ -5,11 +5,12 @@ from typing import List, Optional, Tuple
 import colorful
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption, Supervision
+from lattifai.caption import Supervision
 from lattifai.config import AlignmentConfig
+from lattifai.data import Caption
 from lattifai.utils import safe_print
 
-from .sentence_splitter import END_PUNCTUATION
+from .punctuation import END_PUNCTUATION
 
 
 class Segmenter:

lattifai/alignment/text_align.py

@@ -9,7 +9,8 @@ import regex
 from error_align import error_align
 from error_align.utils import DELIMITERS, NUMERIC_TOKEN, STANDARD_TOKEN, Alignment, OpType
 
-from lattifai.caption import Caption, Supervision
+from lattifai.caption import Supervision
+from lattifai.data import Caption
 from lattifai.utils import safe_print
 
 from .punctuation import PUNCTUATION

lattifai/alignment/tokenizer.py

@@ -6,9 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import numpy as np
 
-# from lattifai.caption import Supervision
-from lhotse.supervision import SupervisionSegment as Supervision  # NOTE: Transcriber SupervisionSegment
-
+from lattifai.caption import SentenceSplitter, Supervision
 from lattifai.caption import normalize_text as normalize_html_text
 from lattifai.errors import (
     LATTICE_DECODING_FAILURE_HELP,
@@ -17,9 +15,7 @@ from lattifai.errors import (
     QuotaExceededError,
 )
 
-from .phonemizer import G2Phonemizer
 from .punctuation import PUNCTUATION, PUNCTUATION_SPACE
-from .sentence_splitter import SentenceSplitter
 from .text_align import TextAlignResult
 
 MAXIMUM_WORD_LENGTH = 40
@@ -174,13 +170,16 @@ class LatticeTokenizer:
         tokenizer.dictionaries = defaultdict(list, data["dictionaries"])
         tokenizer.oov_word = data["oov_word"]
 
+        # Lazy load G2P model only if it exists (avoids PyTorch dependency)
         g2pp_model_path = f"{model_path}/g2pp.bin" if Path(f"{model_path}/g2pp.bin").exists() else None
-        if g2pp_model_path:
-            tokenizer.g2p_model = G2Phonemizer(g2pp_model_path, device=device)
+        g2p_model_path = f"{model_path}/g2p.bin" if Path(f"{model_path}/g2p.bin").exists() else None
+        g2p_path = g2pp_model_path or g2p_model_path
+        if g2p_path:
+            from .phonemizer import G2Phonemizer
+
+            tokenizer.g2p_model = G2Phonemizer(g2p_path, device=device)
         else:
-            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.g2p_model = None
 
         tokenizer.device = device
         tokenizer.add_special_tokens()
@@ -246,9 +245,24 @@ class LatticeTokenizer:
         self.init_sentence_splitter()
         return self.sentence_splitter.split_sentences(supervisions, strip_whitespace=strip_whitespace)
 
+    def _get_client_info(self) -> Dict[str, Optional[str]]:
+        """Get client identification info for usage tracking."""
+        try:
+            from importlib.metadata import version
+
+            return {"client_name": "python-sdk", "client_version": version("lattifai")}
+        except Exception:
+            return {"client_name": "python-sdk", "client_version": "unknown"}
+
     def tokenize(
-        self, supervisions: Union[List[Supervision], TextAlignResult], split_sentence: bool = False, boost: float = 0.0
+        self,
+        supervisions: Union[List[Supervision], TextAlignResult],
+        split_sentence: bool = False,
+        boost: float = 0.0,
+        transition_penalty: Optional[float] = 0.0,
     ) -> Tuple[str, Dict[str, Any]]:
+        client_info = self._get_client_info()
+
         if isinstance(supervisions[0], Supervision):
             if split_sentence:
                 supervisions = self.split_sentences(supervisions)
@@ -260,6 +274,8 @@ class LatticeTokenizer:
                     "model_name": self.model_name,
                     "supervisions": [s.to_dict() for s in supervisions],
                     "pronunciation_dictionaries": pronunciation_dictionaries,
+                    **client_info,
+                    "transition_penalty": transition_penalty,
                 },
             )
         else:
@@ -274,6 +290,7 @@ class LatticeTokenizer:
                     "transcription": [s.to_dict() for s in supervisions[1]],
                     "pronunciation_dictionaries": pronunciation_dictionaries,
                     "boost": boost,
+                    **client_info,
                 },
             )
 
@@ -297,8 +314,10 @@ class LatticeTokenizer:
         return_details: bool = False,
         start_margin: float = 0.08,
         end_margin: float = 0.20,
+        check_sanity: bool = True,
     ) -> List[Supervision]:
-        emission, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
+        emission_stats, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
+        # emission_stats is a dict with 'max_probs' and 'aligned_probs' (unified for batch and streaming)
         if isinstance(supervisions[0], Supervision):
             response = self.client_wrapper.post(
                 "detokenize",
@@ -314,6 +333,7 @@ class LatticeTokenizer:
                     "destroy_lattice": True,
                     "start_margin": start_margin,
                     "end_margin": end_margin,
+                    "check_sanity": check_sanity,
                 },
             )
         else:
@@ -331,6 +351,7 @@ class LatticeTokenizer:
                     "destroy_lattice": True,
                     "start_margin": start_margin,
                     "end_margin": end_margin,
+                    "check_sanity": check_sanity,
                 },
             )
 
@@ -350,9 +371,8 @@ class LatticeTokenizer:
 
         alignments = [Supervision.from_dict(s) for s in result["supervisions"]]
 
-        if emission is not None and return_details:
-            # Add emission confidence scores for segments and word-level alignments
-            _add_confidence_scores(alignments, emission, labels[0], frame_shift, offset)
+        # Add emission confidence scores for segments and word-level alignments
+        _add_confidence_scores(alignments, emission_stats, frame_shift, offset)
 
         if isinstance(supervisions[0], Supervision):
             alignments = _update_alignments_speaker(supervisions, alignments)
@@ -365,8 +385,7 @@ class LatticeTokenizer:
 
 def _add_confidence_scores(
     supervisions: List[Supervision],
-    emission: np.ndarray,
-    labels: List[int],
+    emission_stats: Dict[str, np.ndarray],
     frame_shift: float,
     offset: float = 0.0,
 ) -> None:
@@ -379,29 +398,37 @@ def _add_confidence_scores(
 
     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
+        emission_stats: Dict with 'max_probs' and 'aligned_probs' arrays
         frame_shift: Frame shift in seconds for converting frames to time
+        offset: Time offset in seconds
     """
-    tokens = np.array(labels, dtype=np.int64)
+    max_probs = emission_stats["max_probs"]
+    aligned_probs = emission_stats["aligned_probs"]
+    diffprobs_full = max_probs - aligned_probs
 
     for supervision in supervisions:
         start_frame = int((supervision.start - offset) / frame_shift)
         end_frame = int((supervision.end - offset) / frame_shift)
 
-        # Compute segment-level confidence
-        probabilities = np.exp(emission[0, start_frame:end_frame])
-        aligned = probabilities[range(0, end_frame - start_frame), tokens[start_frame:end_frame]]
-        diffprobs = np.max(probabilities, axis=-1) - aligned
-        supervision.score = round(1.0 - diffprobs.mean(), ndigits=4)
+        # Clamp to valid range
+        start_frame = max(0, min(start_frame, len(diffprobs_full) - 1))
+        end_frame = max(start_frame + 1, min(end_frame, len(diffprobs_full)))
+
+        diffprobs = diffprobs_full[start_frame:end_frame]
+        if len(diffprobs) > 0:
+            supervision.score = round(1.0 - diffprobs.mean().item(), ndigits=4)
 
-        # Compute word-level confidence if alignment exists
+        # Word-level confidence
         if hasattr(supervision, "alignment") and supervision.alignment:
             words = supervision.alignment.get("word", [])
             for w, item in enumerate(words):
-                start = int((item.start - offset) / frame_shift) - start_frame
-                end = int((item.end - offset) / frame_shift) - start_frame
-                words[w] = item._replace(score=round(1.0 - diffprobs[start:end].mean(), ndigits=4))
+                start = int((item.start - offset) / frame_shift)
+                end = int((item.end - offset) / frame_shift)
+                start = max(0, min(start, len(diffprobs_full) - 1))
+                end = max(start + 1, min(end, len(diffprobs_full)))
+                word_diffprobs = diffprobs_full[start:end]
+                if len(word_diffprobs) > 0:
+                    words[w] = item._replace(score=round(1.0 - word_diffprobs.mean().item(), ndigits=4))
 
 
 def _update_alignments_speaker(supervisions: List[Supervision], alignments: List[Supervision]) -> List[Supervision]: