lattifai 0.1.4__py3-none-any.whl

lattifai/tokenizers/tokenizer.py

@@ -0,0 +1,143 @@
+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,
+        g2p_model_path: Optional[str] = None,
+        device: str = 'cpu',
+        compressed: bool = True,
+    ):
+        """Load tokenizer from exported binary file"""
+        if compressed:
+            with gzip.open(model_path, 'rb') as f:
+                data = pickle.load(f)
+        else:
+            with open(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']
+        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)

lattifai/workers/__init__.py

@@ -0,0 +1,3 @@
+from .lattice1_alpha import Lattice1AlphaWorker
+
+__all__ = ['Lattice1AlphaWorker']

lattifai/workers/lattice1_alpha.py

@@ -0,0 +1,119 @@
+import json
+import time
+from collections import defaultdict
+from typing import Any, BinaryIO, Dict, Tuple, Union
+
+import numpy as np
+import onnxruntime as ort
+import torch
+import torchaudio
+from lhotse import FbankConfig
+from lhotse.features.kaldi.layers import Wav2LogFilterBank
+from lhotse.utils import Pathlike
+
+
+class Lattice1AlphaWorker:
+    """Worker for processing audio with LatticeGraph."""
+
+    def __init__(self, model_path: Pathlike, device: str = 'cpu', num_threads: int = 8) -> None:
+        if device != 'cpu':
+            raise NotImplementedError(f'Only cpu is supported for now, got device={device}.')
+        self.config = json.load(open(f'{model_path}/config.json'))
+
+        # SessionOptions
+        sess_options = ort.SessionOptions()
+        # sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        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')
+
+        providers = []
+        if device.startswith('cuda') or ort.get_all_providers().count('CUDAExecutionProvider') > 0:
+            providers.append('CUDAExecutionProvider')
+        self.acoustic_ort = ort.InferenceSession(
+            f'{model_path}/acoustic_opt.onnx',
+            sess_options,
+            providers=providers + ['CoreMLExecutionProvider', 'CPUExecutionProvider'],
+        )
+        config = FbankConfig(num_mel_bins=80, device=device, snip_edges=False)
+        config_dict = config.to_dict()
+        config_dict.pop('device')
+        self.extractor = Wav2LogFilterBank(**config_dict).to(device).eval()
+
+        self.device = torch.device(device)
+        self.timings = defaultdict(lambda: 0.0)
+
+    @torch.inference_mode()
+    def emission(self, audio: torch.Tensor) -> torch.Tensor:
+        _start = time.time()
+        # audio -> features -> emission
+        features = self.extractor(audio)  # (1, T, D)
+        ort_inputs = {
+            'features': features.cpu().numpy(),
+            'feature_lengths': np.array([features.size(1)], dtype=np.int64),
+        }
+        emission = self.acoustic_ort.run(None, ort_inputs)[0]  # (1, T, vocab_size) numpy
+        self.timings['emission'] += time.time() - _start
+        return torch.from_numpy(emission).to(self.device)  # (1, T, vocab_size) torch
+
+    def load_audio(self, audio: Union[Pathlike, BinaryIO]) -> Tuple[torch.Tensor, int]:
+        # load audio
+        waveform, sample_rate = torchaudio.load(audio, channels_first=True)
+        if waveform.size(0) > 1:  # TODO: support choose channel
+            waveform = torch.mean(waveform, dim=0, keepdim=True)
+        if sample_rate != self.config['sample_rate']:
+            waveform = torchaudio.functional.resample(waveform, sample_rate, self.config['sample_rate'])
+        return waveform
+
+    def alignment(
+        self, audio: Union[Union[Pathlike, BinaryIO], torch.tensor], lattice_graph: Tuple[str, int, float]
+    ) -> Dict[str, Any]:
+        """Process audio with LatticeGraph.
+
+        Args:
+            audio: Audio file path or binary data
+            lattice_graph: LatticeGraph data
+
+        Returns:
+            Processed LatticeGraph
+        """
+        # load audio
+        if isinstance(audio, torch.Tensor):
+            waveform = audio
+        else:
+            waveform = self.load_audio(audio)  # (1, L)
+
+        _start = time.time()
+        emission = self.emission(waveform.to(self.device))  # (1, T, vocab_size)
+        self.timings['emission'] += time.time() - _start
+
+        import k2
+        from lattifai_core.lattice.decode import align_segments
+
+        lattice_graph_str, final_state, acoustic_scale = lattice_graph
+
+        _start = time.time()
+        # graph
+        decoding_graph = k2.Fsa.from_str(lattice_graph_str, acceptor=False)
+        decoding_graph.requires_grad_(False)
+        decoding_graph = k2.arc_sort(decoding_graph)
+        decoding_graph.skip_id = int(final_state)
+        decoding_graph.return_id = int(final_state + 1)
+        self.timings['decoding_graph'] += time.time() - _start
+
+        _start = time.time()
+        results, labels = align_segments(
+            emission.to(self.device) * acoustic_scale,
+            decoding_graph.to(self.device),
+            torch.tensor([emission.shape[1]], dtype=torch.int32),
+            search_beam=100,
+            output_beam=40,
+            min_active_states=200,
+            max_active_states=10000,
+            subsampling_factor=1,
+            reject_low_confidence=False,
+        )
+        self.timings['align_segments'] += time.time() - _start
+
+        channel = 0
+        return emission, results, labels, 0.02, 0.0, channel  # frame_shift=20ms, offset=0.0s