# Frozen inference function extraction. See source provenance and licenses. import gc import math import re from collections import Counter from itertools import groupby import numpy as np import librosa import torch from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq EPS = 1e-8 EPS_TOKEN = '' def normalize_phone(phone): phone = phone.upper().strip() phone = phone.replace('[', '').replace(']', '') phone = re.sub(r'[_\d].*', '', phone) return phone def normalize_word(word): word = word.upper().strip() word = word.replace('[', '').replace(']', '') word = re.sub(r"[^A-Z']", '', word) return word def monotonic_align(neg_log_probs, phone_ids): frame_num = neg_log_probs.shape[0] token_num = len(phone_ids) if frame_num < token_num: raise ValueError(f'Not enough frames ({frame_num}) to align {token_num} phones.') dp = np.full((frame_num, token_num), np.inf, dtype=np.float32) back = np.zeros((frame_num, token_num), dtype=np.int8) dp[0, 0] = neg_log_probs[0, phone_ids[0]] for t in range(1, frame_num): max_token = min(t, token_num - 1) for tok in range(max_token + 1): stay_cost = dp[t - 1, tok] move_cost = dp[t - 1, tok - 1] if tok > 0 else np.inf if move_cost < stay_cost: dp[t, tok] = move_cost + neg_log_probs[t, phone_ids[tok]] back[t, tok] = 1 else: dp[t, tok] = stay_cost + neg_log_probs[t, phone_ids[tok]] path = np.zeros(frame_num, dtype=np.int32) tok = token_num - 1 for t in range(frame_num - 1, -1, -1): path[t] = tok if t > 0 and back[t, tok] == 1: tok -= 1 return path def select_visible_frames(probs, keep_mask, audio_end, frame_step): frame_limit = max(int(math.ceil(audio_end / max(frame_step, EPS))), 1) frame_mask = np.zeros(len(probs), dtype=bool) frame_mask[: min(frame_limit, len(probs))] = True visible_mask = keep_mask & frame_mask return np.flatnonzero(visible_mask), probs[visible_mask] def softmax_1d(values): arr = np.asarray(values, dtype=np.float64) if arr.size == 0: return arr.astype(np.float32) pos_inf = np.isposinf(arr) if np.any(pos_inf): out = np.zeros_like(arr, dtype=np.float64) out[pos_inf] = 1.0 / float(np.sum(pos_inf)) return out.astype(np.float32) finite = np.isfinite(arr) if not np.any(finite): return np.full(arr.shape, 1.0 / float(arr.size), dtype=np.float32) min_finite = np.min(arr[finite]) arr = np.where(finite, arr, min_finite - 50.0) arr = arr - np.max(arr) exp_arr = np.exp(arr) return (exp_arr / np.clip(exp_arr.sum(), EPS, None)).astype(np.float32) def entropy_np(probs): probs = np.asarray(probs, dtype=np.float64) probs = probs[probs > 0] if probs.size == 0: return 0.0 return float(-np.sum(probs * np.log(np.clip(probs, EPS, None)))) def top_margin(probs): probs = np.asarray(probs, dtype=np.float32) if probs.size == 0: return 0.0, 0.0 top = np.sort(probs)[::-1] top1 = float(top[0]) top2 = float(top[1]) if top.size > 1 else 0.0 return top1, float(top1 - top2) def js_divergence(p, q): p = np.asarray(p, dtype=np.float64) q = np.asarray(q, dtype=np.float64) p = p / np.clip(p.sum(), EPS, None) q = q / np.clip(q.sum(), EPS, None) m = 0.5 * (p + q) kl_pm = np.sum(p * (np.log(np.clip(p, EPS, None)) - np.log(np.clip(m, EPS, None)))) kl_qm = np.sum(q * (np.log(np.clip(q, EPS, None)) - np.log(np.clip(m, EPS, None)))) return float(0.5 * (kl_pm + kl_qm)) def safe_json(value): if isinstance(value, dict): return {str(key): safe_json(val) for key, val in value.items()} if isinstance(value, (list, tuple)): return [safe_json(item) for item in value] if isinstance(value, np.ndarray): return safe_json(value.tolist()) if isinstance(value, (np.floating, np.integer)): return value.item() return value def softmax_np_matrix(values): arr = np.asarray(values, dtype=np.float64) if arr.size == 0: return arr.astype(np.float32) arr = arr - np.max(arr, axis=-1, keepdims=True) exp_arr = np.exp(arr) return (exp_arr / np.clip(np.sum(exp_arr, axis=-1, keepdims=True), EPS, None)).astype(np.float32) def audio_array_logits(audio, processor, model, sample_rate, device): audio = np.asarray(audio, dtype=np.float32) if hasattr(processor, 'audio_preprocess'): inputs = processor.audio_preprocess(audio, sr=sample_rate) inputs = torch.tensor(inputs, dtype=torch.float32).unsqueeze(0).to(device) model_inputs = {'input_values': inputs} else: model_inputs = processor(audio, sampling_rate=sample_rate, return_tensors='pt') model_inputs = {key: value.to(device) for key, value in model_inputs.items()} with torch.no_grad(): logits = model(**model_inputs).logits[0].detach().cpu().numpy() return softmax_np_matrix(logits), float(len(audio) / sample_rate) def non_overlapping_occurrences(tokens, pattern): positions = [] pattern_size = len(pattern) cursor = 0 while cursor <= len(tokens) - pattern_size: if tokens[cursor : cursor + pattern_size] == pattern: positions.append(cursor) cursor += pattern_size else: cursor += 1 return positions def find_repeated_ngram( words, min_repeats, max_ngram_size=12, strict=False, min_coverage=0.6, dominant_token_ratio=0.5, ): if min_repeats <= 0: return None tokens = [str(word['text']) for word in words] if not tokens: return None for ngram_size in range(1, min(max_ngram_size, len(tokens)) + 1): span = ngram_size * min_repeats for start in range(0, len(tokens) - span + 1): pattern = tokens[start : start + ngram_size] repeats = 1 cursor = start + ngram_size while tokens[cursor : cursor + ngram_size] == pattern: repeats += 1 cursor += ngram_size if repeats >= min_repeats: return { 'start': int(start), 'ngram_size': int(ngram_size), 'repeats': int(repeats), 'pattern': pattern, 'detection': 'contiguous_cycle', } if not strict: return None token_counts = Counter(tokens) dominant_token, dominant_count = token_counts.most_common(1)[0] token_ratio = float(dominant_count) / float(len(tokens)) if dominant_count >= max(min_repeats * 2, 8) and token_ratio >= dominant_token_ratio: return { 'start': None, 'ngram_size': 1, 'repeats': int(dominant_count), 'pattern': [dominant_token], 'coverage': token_ratio, 'detection': 'dominant_token', } max_size = min(max_ngram_size, len(tokens) // max(min_repeats, 1)) for ngram_size in range(max_size, 1, -1): candidates = Counter( tuple(tokens[start : start + ngram_size]) for start in range(0, len(tokens) - ngram_size + 1) ) for pattern, raw_count in candidates.most_common(): if raw_count < min_repeats: break positions = non_overlapping_occurrences(tokens, list(pattern)) if len(positions) < min_repeats: continue coverage = float(len(positions) * ngram_size) / float(len(tokens)) if coverage >= min_coverage: return { 'start': int(positions[0]), 'ngram_size': int(ngram_size), 'repeats': int(len(positions)), 'pattern': list(pattern), 'coverage': coverage, 'detection': 'high_coverage_periodicity', } return None class PhoneMapper: def __init__(self, lexicon, phn_dict): self.lexicon = lexicon self.phn_dict = phn_dict self.cache = {} self.fallback_phone = 'AH' if 'AH' in phn_dict else (next(iter(phn_dict)) if phn_dict else None) self.g2p = None try: from g2p_en import G2p self.g2p = G2p() except Exception: pass def word_to_phones(self, word): word = normalize_word(word) if word in self.lexicon: return list(self.lexicon[word]), 'lexicon' if word in self.cache: return self.cache[word] if self.g2p is not None: try: phones = [] for raw_phone in self.g2p(word.lower()): phone = normalize_phone(raw_phone) if phone and phone in self.phn_dict: phones.append(phone) if phones: result = (phones, 'g2p_en') self.cache[word] = result return result except Exception: self.g2p = None if self.fallback_phone is None: result = ([], 'unmapped') else: approx_len = max(1, int(math.ceil(len(re.sub(r'[^A-Z]', '', word)) / 3.0))) result = ([self.fallback_phone] * approx_len, 'fallback_phone') self.cache[word] = result return result def words_to_phone_sequence(self, words): phones = [] phone_to_word = [] source_counts = Counter() for word_idx, word in enumerate(words): cur_phones, source = self.word_to_phones(word) source_counts[source] += 1 for phone in cur_phones: phones.append(phone) phone_to_word.append(word_idx) return phones, phone_to_word, dict(source_counts) class WhisperNBestGenerator: def __init__(self, model_name, language, device, nbest, beam_size, max_new_tokens, no_repeat_ngram_size): self.processor = AutoProcessor.from_pretrained(model_name) torch_dtype = torch.float16 if str(device).startswith('cuda') else torch.float32 self.model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name, torch_dtype=torch_dtype) self.model.to(device) self.model.eval() self.model_dtype = next(self.model.parameters()).dtype self.language = language self.device = device self.nbest = int(nbest) self.beam_size = max(int(beam_size), int(nbest)) self.max_new_tokens = int(max_new_tokens) self.no_repeat_ngram_size = int(no_repeat_ngram_size) self.forced_decoder_ids = None if hasattr(self.processor, 'get_decoder_prompt_ids'): try: self.forced_decoder_ids = self.processor.get_decoder_prompt_ids( language=language, task='transcribe', ) except Exception: self.forced_decoder_ids = None @torch.inference_mode() def generate(self, audio, sample_rate, audio_end): inputs = self.processor( np.asarray(audio, dtype=np.float32), sampling_rate=sample_rate, return_tensors='pt', return_attention_mask=True, ) inputs = {key: value.to(self.device) for key, value in inputs.items()} for key in ('input_features', 'input_values'): if key in inputs and torch.is_floating_point(inputs[key]): inputs[key] = inputs[key].to(dtype=self.model_dtype) generate_kwargs = { 'num_beams': self.beam_size, 'num_return_sequences': self.nbest, 'max_new_tokens': self.max_new_tokens, 'return_dict_in_generate': True, 'output_scores': True, 'do_sample': False, } if self.no_repeat_ngram_size > 0: generate_kwargs['no_repeat_ngram_size'] = self.no_repeat_ngram_size if self.forced_decoder_ids is not None: generate_kwargs['forced_decoder_ids'] = self.forced_decoder_ids output = self.model.generate(**inputs, **generate_kwargs) if torch.is_tensor(output): sequences = output sequence_scores = [0.0 for _ in range(int(sequences.shape[0]))] transition_scores = [] else: sequences = output.sequences if getattr(output, 'sequences_scores', None) is not None: sequence_scores = output.sequences_scores.detach().float().cpu().numpy().tolist() else: sequence_scores = [0.0 for _ in range(int(sequences.shape[0]))] transition_scores = self._transition_scores(output) texts = self.processor.batch_decode(sequences, skip_special_tokens=True) rows = [] for rank, (text, sequence_score) in enumerate(zip(texts, sequence_scores)): token_ids = sequences[rank].detach().cpu().tolist() token_logprobs = ( transition_scores[rank] if rank < len(transition_scores) else [float(sequence_score) for _ in token_ids] ) token_rows, word_rows = self._token_word_scores(token_ids, token_logprobs) words = [row['word'] for row in word_rows] length_norm = ( float(np.mean([row['logprob'] for row in token_rows])) if token_rows else float(sequence_score) ) rows.append({ 'rank': int(rank), 'text': ' '.join(words).lower(), 'words': words, 'logprob': float(length_norm), 'sequence_score': float(sequence_score), 'length_normalized_sequence_score': float(length_norm), 'token_ids': [int(row['token_id']) for row in token_rows], 'token_logprobs': [float(row['logprob']) for row in token_rows], 'token_confidences': [float(row['confidence']) for row in token_rows], 'word_token_ranges': [[int(row['token_start']), int(row['token_end'])] for row in word_rows], 'word_logprobs': [float(row['logprob']) for row in word_rows], 'word_confidences': [float(row['confidence']) for row in word_rows], 'word_timestamps': estimate_word_timestamps(words, audio_end), }) if str(self.device).startswith('cuda'): torch.cuda.empty_cache() gc.collect() return rows def _transition_scores(self, output): if getattr(output, 'scores', None) is None: return [] try: sequences = output.sequences.detach().cpu() step_scores = [score.detach().float().cpu() for score in output.scores] if not step_scores: return [] beam_indices = getattr(output, 'beam_indices', None) if beam_indices is not None: beam_indices = beam_indices.detach().cpu() generated_steps = len(step_scores) start = max(0, int(sequences.shape[1]) - generated_steps) rows = [] for seq_idx in range(int(sequences.shape[0])): cur = [] for step, logits in enumerate(step_scores): token_pos = start + step if token_pos >= int(sequences.shape[1]): break token_id = int(sequences[seq_idx, token_pos]) if token_id < 0 or token_id >= int(logits.shape[-1]): cur.append(0.0) continue score_row = min(seq_idx, int(logits.shape[0]) - 1) if beam_indices is not None and seq_idx < int(beam_indices.shape[0]) and token_pos < int(beam_indices.shape[1]): candidate = int(beam_indices[seq_idx, token_pos]) if 0 <= candidate < int(logits.shape[0]): score_row = candidate log_probs = torch.log_softmax(logits[score_row], dim=-1) cur.append(float(log_probs[token_id].item())) rows.append(cur) return rows except Exception: return [] def _token_word_scores(self, token_ids, token_logprobs): tokenizer = self.processor.tokenizer special_ids = set(getattr(tokenizer, 'all_special_ids', []) or []) if len(token_logprobs) != len(token_ids): aligned = [0.0 for _ in token_ids] offset = max(0, len(token_ids) - len(token_logprobs)) for idx, value in enumerate(token_logprobs[-len(token_ids) :]): pos = offset + idx if pos < len(aligned): aligned[pos] = float(value) token_logprobs = aligned token_rows = [] words = [] current = [] current_token_ids = [] current_logprobs = [] current_start = None def flush(): nonlocal current, current_token_ids, current_logprobs, current_start word = normalize_word(''.join(current)) if word: avg_logprob = float(np.mean(current_logprobs)) if current_logprobs else 0.0 words.append({ 'word': word, 'token_start': int(current_start if current_start is not None else 0), 'token_end': int((current_token_ids[-1] + 1) if current_token_ids else 0), 'logprob': avg_logprob, 'confidence': float(np.clip(np.exp(avg_logprob), 0.0, 1.0)), }) current = [] current_token_ids = [] current_logprobs = [] current_start = None for pos, token_id in enumerate(token_ids): if int(token_id) in special_ids: continue logprob = float(token_logprobs[pos]) if pos < len(token_logprobs) else 0.0 try: text = tokenizer.decode([int(token_id)], skip_special_tokens=True, clean_up_tokenization_spaces=False) except TypeError: text = tokenizer.decode([int(token_id)], skip_special_tokens=True) if not text: continue token_rows.append({ 'token_id': int(token_id), 'logprob': logprob, 'confidence': float(np.clip(np.exp(logprob), 0.0, 1.0)), }) for char in text: if char.isalpha() or char == "'": if current_start is None: current_start = pos current.append(char) current_token_ids.append(pos) current_logprobs.append(logprob) else: flush() flush() return token_rows, words def normalize_text_to_words(text): tokens = [] for token in re.findall(r"[A-Za-z']+", text): norm = normalize_word(token) if norm: tokens.append(norm) return tokens def estimate_word_timestamps(words, audio_end): if not words: return [] step = float(audio_end) / max(len(words), 1) rows = [] for idx, word in enumerate(words): rows.append({ 'word': word.lower(), 'start': float(idx * step), 'end': float(min(audio_end, (idx + 1) * step)), 'source': 'duration_proportional_estimate', }) return rows def hypothesis_phone_rows(row, phone_mapper, phn_dict): phones = [] phone_ids = [] phone_to_word = [] phone_confidences = [] source_counts = Counter() word_confidences = row.get('word_confidences') or [1.0 for _ in row.get('words', [])] for word_idx, word in enumerate(row.get('words', [])): cur_phones, source = phone_mapper.word_to_phones(word) source_counts[source] += 1 word_conf = float(word_confidences[word_idx]) if word_idx < len(word_confidences) else 1.0 word_conf = float(np.clip(word_conf, 0.0, 1.0)) for phone in cur_phones: if phone not in phn_dict: continue phones.append(phone) phone_ids.append(int(phn_dict[phone])) phone_to_word.append(int(word_idx)) phone_confidences.append(word_conf) return { 'phones': phones, 'phone_ids': phone_ids, 'phone_to_word': phone_to_word, 'phone_confidences': phone_confidences, 'source_counts': dict(source_counts), } def align_hypothesis_with_charsiu(item, row, phone_mapper, phn_dict, phone_to_frame_id, audio_end): hyp = hypothesis_phone_rows(row, phone_mapper, phn_dict) words = row.get('words', []) if not words or not hyp['phone_ids']: row['word_timestamps'] = [] row['timestamp_source'] = 'empty_hypothesis' row['phone_times'] = [] row['phone_acoustic_supports'] = [] return row kept_indices, kept_probs = select_visible_frames(item['probs'], item['keep_mask'], audio_end, item['frame_step']) frame_phone_ids = [] id_to_phone = {idx: phone for phone, idx in phn_dict.items()} try: for phone_idx in hyp['phone_ids']: phone = id_to_phone[int(phone_idx)] frame_phone_ids.append(phone_to_frame_id[phone]) if kept_probs.shape[0] < len(frame_phone_ids): raise ValueError('not_enough_visible_frames_for_hypothesis') path = monotonic_align(-np.log(np.clip(kept_probs, EPS, None)), frame_phone_ids) phone_times = [] phone_acoustic_supports = [] for phone_pos, frame_phone_id in enumerate(frame_phone_ids): local_frames = kept_indices[path == phone_pos] if local_frames.size == 0: raise ValueError('empty_aligned_phone_span') start_time = float(local_frames.min() * item['frame_step']) end_time = float((local_frames.max() + 1) * item['frame_step']) support = float(np.mean(item['probs'][local_frames, frame_phone_id])) phone_times.append((start_time, min(float(audio_end), end_time))) phone_acoustic_supports.append(float(np.clip(support, 0.0, 1.0))) word_rows = [] for word_idx, word in enumerate(words): positions = [idx for idx, cur_word in enumerate(hyp['phone_to_word']) if cur_word == word_idx] if not positions: continue start = min(phone_times[idx][0] for idx in positions) end = max(phone_times[idx][1] for idx in positions) word_rows.append({ 'word': word, 'start': float(start), 'end': float(end), 'source': 'charsiu_hypothesis_alignment', 'word_logprob': float(row.get('word_logprobs', [0.0] * len(words))[word_idx]) if word_idx < len(row.get('word_logprobs', [])) else 0.0, 'word_confidence': float(row.get('word_confidences', [1.0] * len(words))[word_idx]) if word_idx < len(row.get('word_confidences', [])) else 1.0, }) row['word_timestamps'] = word_rows row['timestamp_source'] = 'charsiu_hypothesis_alignment' row['phone_times'] = phone_times row['phone_acoustic_supports'] = phone_acoustic_supports except Exception as exc: row['word_timestamps'] = estimate_word_timestamps(words, audio_end) row['timestamp_source'] = 'duration_proportional_fallback' row['timestamp_fallback_reason'] = repr(exc) phone_times = [] phone_acoustic_supports = [] for phone_idx, word_idx in enumerate(hyp['phone_to_word']): if word_idx < len(row['word_timestamps']): word_time = row['word_timestamps'][word_idx] phone_times.append((float(word_time['start']), float(word_time['end']))) else: phone_times.append((0.0, float(audio_end))) phone_acoustic_supports.append(0.0) row['phone_times'] = phone_times row['phone_acoustic_supports'] = phone_acoustic_supports row.update(hyp) return row def interval_overlap_ratio(a, b): if a is None or b is None: return 0.0 a0, a1 = a b0, b1 = b inter = max(0.0, min(float(a1), float(b1)) - max(float(a0), float(b0))) union = max(float(a1), float(b1)) - min(float(a0), float(b0)) if union <= 0: return 0.0 return inter / union def align_token_sequences( ref_tokens, hyp_tokens, mismatch_cost=1.0, ref_times=None, hyp_times=None, hyp_confidences=None, time_overlap_reward=0.0, ): n = len(ref_tokens) m = len(hyp_tokens) dp = np.zeros((n + 1, m + 1), dtype=np.float32) back = {} for i in range(1, n + 1): dp[i, 0] = dp[i - 1, 0] + 1.0 back[(i, 0)] = (i - 1, 0, 'delete') for j in range(1, m + 1): dp[0, j] = dp[0, j - 1] + 1.0 back[(0, j)] = (0, j - 1, 'insert') for i in range(1, n + 1): for j in range(1, m + 1): sub_cost = 0.0 if ref_tokens[i - 1] == hyp_tokens[j - 1] else mismatch_cost if ref_times is not None and hyp_times is not None: sub_cost -= float(time_overlap_reward) * interval_overlap_ratio(ref_times[i - 1], hyp_times[j - 1]) if hyp_confidences is not None: sub_cost -= 0.1 * float(np.clip(hyp_confidences[j - 1], 0.0, 1.0)) sub_cost = max(0.0, sub_cost) candidates = [ (dp[i - 1, j - 1] + sub_cost, i - 1, j - 1, 'match'), (dp[i - 1, j] + 1.0, i - 1, j, 'delete'), (dp[i, j - 1] + 1.0, i, j - 1, 'insert'), ] cost, prev_i, prev_j, op = min(candidates, key=lambda item: item[0]) dp[i, j] = cost back[(i, j)] = (prev_i, prev_j, op) pairs = [] i, j = n, m while i > 0 or j > 0: prev_i, prev_j, op = back[(i, j)] if op == 'match': pairs.append((i - 1, j - 1)) elif op == 'delete': pairs.append((i - 1, None)) else: pairs.append((None, j - 1)) i, j = prev_i, prev_j pairs.reverse() return pairs def slot_top_phone(slot_counts, eps_index): best_idx = None best_mass = -1.0 for phone_idx, mass in slot_counts.items(): if phone_idx == eps_index: continue if mass > best_mass: best_idx = phone_idx best_mass = mass return best_idx def slot_time(slot): mass = max(float(slot.get('time_mass', 0.0)), EPS) if mass <= EPS: return None return float(slot.get('start_sum', 0.0) / mass), float(slot.get('end_sum', 0.0) / mass) def add_slot_phone_mass(slot, phone_idx, phone_mass, eps_index, eps_mass=0.0, phone_time=None, confidence=1.0): if phone_mass > 0: slot['counts'][int(phone_idx)] += float(phone_mass) if eps_mass > 0: slot['counts'][eps_index] += float(eps_mass) if phone_time is not None and phone_mass > 0: slot['start_sum'] = slot.get('start_sum', 0.0) + float(phone_time[0]) * float(phone_mass) slot['end_sum'] = slot.get('end_sum', 0.0) + float(phone_time[1]) * float(phone_mass) slot['time_mass'] = slot.get('time_mass', 0.0) + float(phone_mass) slot['confidence_sum'] = slot.get('confidence_sum', 0.0) + float(confidence) * max(float(phone_mass), 0.0) slot['confidence_mass'] = slot.get('confidence_mass', 0.0) + max(float(phone_mass), 0.0) def insert_slot(slots, insert_at, eps_index, prior_mass): slot = {'counts': Counter(), 'start_sum': 0.0, 'end_sum': 0.0, 'time_mass': 0.0} if prior_mass > 0: slot['counts'][eps_index] += float(prior_mass) slots.insert(insert_at, slot) return slot def build_pcn_from_hypotheses(hypotheses, phone_mapper, phn_dict): eps_index = len(phn_dict) phone_dim = len(phn_dict) + 1 hyp_logprobs = [float(row.get('length_normalized_sequence_score', row.get('logprob', 0.0))) for row in hypotheses] hyp_weights = softmax_1d(hyp_logprobs) id_to_phone = {idx: phone for phone, idx in phn_dict.items()} hyp_phone_rows = [] for row, weight in zip(hypotheses, hyp_weights): if 'phone_ids' not in row: row = align_hypothesis_with_charsiu( item={'probs': np.zeros((0, phone_dim), dtype=np.float32), 'keep_mask': np.zeros((0,), dtype=np.int32), 'frame_step': 0.02}, row=row, phone_mapper=phone_mapper, phn_dict=phn_dict, phone_to_frame_id={}, audio_end=0.0, ) phone_ids = [int(phone_idx) for phone_idx in row.get('phone_ids', [])] phone_confidences = [float(x) for x in row.get('phone_confidences', [1.0] * len(phone_ids))] phone_acoustic = [float(x) for x in row.get('phone_acoustic_supports', [0.0] * len(phone_ids))] phone_times = row.get('phone_times', [None] * len(phone_ids)) hyp_phone_rows.append({ 'phones': row.get('phones', []), 'phone_ids': phone_ids, 'phone_to_word': list(row.get('phone_to_word', [])), 'phone_confidences': phone_confidences, 'phone_acoustic_supports': phone_acoustic, 'phone_times': phone_times, 'weight': float(weight), 'source_counts': row.get('source_counts', {}), }) slots = [] total_mass = 0.0 for hyp_idx, hyp in enumerate(hyp_phone_rows): phone_ids = hyp['phone_ids'] weight = float(hyp['weight']) if hyp_idx == 0: for local_idx, phone_idx in enumerate(phone_ids): slot = {'counts': Counter(), 'start_sum': 0.0, 'end_sum': 0.0, 'time_mass': 0.0} quality = float(np.clip( 0.5 * hyp['phone_confidences'][local_idx] + 0.5 * hyp['phone_acoustic_supports'][local_idx], 0.0, 1.0, )) add_slot_phone_mass( slot, phone_idx, weight * quality, eps_index, eps_mass=weight * (1.0 - quality), phone_time=hyp['phone_times'][local_idx] if local_idx < len(hyp['phone_times']) else None, confidence=hyp['phone_confidences'][local_idx], ) slots.append(slot) total_mass += weight continue consensus_slot_ids = [] consensus_phone_ids = [] for slot_idx, slot in enumerate(slots): top_idx = slot_top_phone(slot['counts'], eps_index) if top_idx is not None: consensus_slot_ids.append(slot_idx) consensus_phone_ids.append(top_idx) consensus_times = [slot_time(slots[slot_idx]) for slot_idx in consensus_slot_ids] pairs = align_token_sequences( consensus_phone_ids, phone_ids, ref_times=consensus_times, hyp_times=hyp.get('phone_times'), hyp_confidences=hyp.get('phone_confidences'), time_overlap_reward=0.3, ) cursor = 0 matched_slots = set() for ref_pos, hyp_pos in pairs: if ref_pos is None: insert_at = consensus_slot_ids[cursor] if cursor < len(consensus_slot_ids) else len(slots) slot = insert_slot(slots, insert_at, eps_index, total_mass) quality = float(np.clip( 0.5 * hyp['phone_confidences'][hyp_pos] + 0.5 * hyp['phone_acoustic_supports'][hyp_pos], 0.0, 1.0, )) add_slot_phone_mass( slot, phone_ids[hyp_pos], weight * quality, eps_index, eps_mass=weight * (1.0 - quality), phone_time=hyp['phone_times'][hyp_pos] if hyp_pos < len(hyp['phone_times']) else None, confidence=hyp['phone_confidences'][hyp_pos], ) cursor += 1 else: slot_idx = consensus_slot_ids[ref_pos] matched_slots.add(slot_idx) cursor = ref_pos + 1 if hyp_pos is None: slots[slot_idx]['counts'][eps_index] += weight else: quality = float(np.clip( 0.5 * hyp['phone_confidences'][hyp_pos] + 0.5 * hyp['phone_acoustic_supports'][hyp_pos], 0.0, 1.0, )) add_slot_phone_mass( slots[slot_idx], phone_ids[hyp_pos], weight * quality, eps_index, eps_mass=weight * (1.0 - quality), phone_time=hyp['phone_times'][hyp_pos] if hyp_pos < len(hyp['phone_times']) else None, confidence=hyp['phone_confidences'][hyp_pos], ) for slot_idx, slot in enumerate(slots): if slot_idx not in matched_slots and sum(slot['counts'].values()) < total_mass + weight - EPS: if slot_top_phone(slot['counts'], eps_index) is not None: slot['counts'][eps_index] += weight total_mass += weight if not slots: slot = {'counts': Counter()} slot['counts'][eps_index] = 1.0 slots = [slot] total_mass = 1.0 cn_post = np.zeros((len(slots), phone_dim), dtype=np.float32) top_phone_ids = [] slot_times = [] slot_confidences = [] for slot_idx, slot in enumerate(slots): slot_mass = sum(slot['counts'].values()) if slot_mass < total_mass: slot['counts'][eps_index] += total_mass - slot_mass for phone_idx, mass in slot['counts'].items(): cn_post[slot_idx, int(phone_idx)] = float(mass) / max(float(total_mass), EPS) cn_post[slot_idx] /= np.clip(cn_post[slot_idx].sum(), EPS, None) top_phone_ids.append(slot_top_phone(slot['counts'], eps_index)) slot_times.append(slot_time(slot)) conf_mass = max(float(slot.get('confidence_mass', 0.0)), EPS) slot_confidences.append(float(np.clip(slot.get('confidence_sum', 0.0) / conf_mass, 0.0, 1.0))) return { 'cn_post': cn_post, 'top_phone_ids': top_phone_ids, 'slot_times': slot_times, 'slot_confidences': slot_confidences, 'hyp_phone_rows': hyp_phone_rows, 'hyp_weights': hyp_weights.tolist(), 'id_to_phone': id_to_phone, 'eps_index': eps_index, } def top_hyp_word_ids_for_slots(pcn): slot_word_ids = np.zeros((len(pcn['top_phone_ids']),), dtype=np.int32) - 1 if not pcn['hyp_phone_rows']: return slot_word_ids top_hyp = pcn['hyp_phone_rows'][0] hyp_phone_ids = list(top_hyp['phone_ids']) hyp_phone_to_word = list(top_hyp['phone_to_word']) consensus_slot_ids = [] consensus_phone_ids = [] for slot_idx, phone_idx in enumerate(pcn['top_phone_ids']): if phone_idx is not None: consensus_slot_ids.append(slot_idx) consensus_phone_ids.append(phone_idx) pairs = align_token_sequences(consensus_phone_ids, hyp_phone_ids) for ref_pos, hyp_pos in pairs: if ref_pos is None or hyp_pos is None: continue slot_idx = consensus_slot_ids[ref_pos] if hyp_pos < len(hyp_phone_to_word): slot_word_ids[slot_idx] = int(hyp_phone_to_word[hyp_pos]) return slot_word_ids def pcn_stats(cn_post, top_phone_ids, prev_top_phone_ids, eps_index): stats = np.zeros((cn_post.shape[0], 5), dtype=np.float32) stable_prefix = 0 for idx, phone_idx in enumerate(top_phone_ids): if idx < len(prev_top_phone_ids) and phone_idx == prev_top_phone_ids[idx]: stable_prefix += 1 else: break prefix_stability = float(stable_prefix) / float(max(len(top_phone_ids), 1)) for idx, probs in enumerate(cn_post): non_eps = np.delete(probs, eps_index) top1, margin = top_margin(non_eps) stats[idx] = np.array( [ float(probs[eps_index]), entropy_np(probs), top1, margin, prefix_stability, ], dtype=np.float32, ) return stats, prefix_stability def validate_pcn(pcn): cn_post = np.asarray(pcn['cn_post'], dtype=np.float32) if not np.isfinite(cn_post).all(): raise ValueError('pcn_contains_nan_or_inf') sums = cn_post.sum(axis=-1) if not np.allclose(sums, 1.0, atol=1e-3): raise ValueError(f'pcn_posterior_not_normalized:max_abs_err={float(np.max(np.abs(sums - 1.0)))}') eps_index = int(pcn['eps_index']) if eps_index < 0 or eps_index >= cn_post.shape[-1]: raise ValueError('invalid_epsilon_index') eps_probs = cn_post[:, eps_index] if np.any(eps_probs < -1e-6) or np.any(eps_probs > 1.0 + 1e-6): raise ValueError('epsilon_probability_out_of_range') for row in pcn.get('hyp_phone_rows', []): for value in row.get('phone_confidences', []): if not np.isfinite(value) or value < -1e-6 or value > 1.0 + 1e-6: raise ValueError('phone_confidence_out_of_range') for value in pcn.get('slot_confidences', []): if not np.isfinite(value) or value < -1e-6 or value > 1.0 + 1e-6: raise ValueError('slot_confidence_out_of_range') def cumulative_candidate_mask_from_asr(pcn_word_id, word_timestamps, commit_time, audio_end, is_final): mask = np.zeros((pcn_word_id.shape[0],), dtype=np.float32) committed_words = set() for word_idx, word_time in enumerate(word_timestamps or []): end_time = float(word_time.get('end', 0.0)) if end_time <= commit_time + 1e-6 or (is_final and end_time <= audio_end + 1e-6): committed_words.add(int(word_idx)) for slot_idx, word_idx in enumerate(pcn_word_id.tolist()): if int(word_idx) in committed_words: mask[slot_idx] = 1.0 return mask def align_previous_commits(prev_state, current_top_phone_ids): if prev_state is None: return {}, [], [] prev_top_phone_ids = prev_state.get('top_phone_ids', []) prev_committed = set(int(x) for x in prev_state.get('committed_slots', [])) pairs = align_token_sequences(prev_top_phone_ids, current_top_phone_ids, mismatch_cost=1.0) prev_to_current = {} for prev_idx, cur_idx in pairs: if prev_idx is None or cur_idx is None: continue if prev_idx < len(prev_top_phone_ids) and cur_idx < len(current_top_phone_ids): if prev_top_phone_ids[prev_idx] == current_top_phone_ids[cur_idx]: prev_to_current[int(prev_idx)] = int(cur_idx) mapped_old = sorted(prev_to_current[idx] for idx in prev_committed if idx in prev_to_current) dropped = sorted(idx for idx in prev_committed if idx not in prev_to_current) return prev_to_current, mapped_old, dropped def _top_word_timestamps(word_timestamps): rows = word_timestamps or [] if rows and isinstance(rows[0], list): rows = rows[0] return rows if isinstance(rows, list) else [] def _commit_word_metadata(word_timestamps, word_idx): rows = _top_word_timestamps(word_timestamps) row = rows[word_idx] if 0 <= int(word_idx) < len(rows) and isinstance(rows[word_idx], dict) else {} text = normalize_word(str(row.get('word', ''))).casefold() try: start = float(row.get('start')) end = float(row.get('end')) except (TypeError, ValueError): start = None end = None return {'text': text, 'start': start, 'end': end} def _commit_time_iou(left, right): if any(item.get(key) is None for item in (left, right) for key in ('start', 'end')): return 0.0 overlap = max(0.0, min(left['end'], right['end']) - max(left['start'], right['start'])) union = max(left['end'], right['end']) - min(left['start'], right['start']) return overlap / union if union > 0 else 0.0 def _word_level_lineage(prev_state, pcn_word_id, current_to_prev, word_timestamps): """Build deterministic one-to-one predecessor decisions without using supervision.""" current_positions = { word_idx: np.flatnonzero(pcn_word_id == word_idx) for word_idx in sorted(set(int(value) for value in pcn_word_id.tolist() if int(value) >= 0)) } if prev_state is None: return current_positions, { word_idx: { 'candidate_prev_word_id': None, 'selected_prev_word_id': None, 'predecessor_word_ids': [], 'vote_counts': {}, 'mapping_status': 'first_seen', 'split_competitor_count': 0, 'reliable': True, 'text_equal': False, 'time_iou': 0.0, } for word_idx in current_positions } prev_word_ids = np.asarray(prev_state.get('pcn_word_id', []), dtype=np.int32) prev_word_states = prev_state.get('word_states', {}) decisions = {} competitors_by_prev = {} for word_idx, positions in current_positions.items(): votes = Counter() for position in positions.tolist(): old_slot = int(current_to_prev[position]) if position < len(current_to_prev) else -1 if 0 <= old_slot < len(prev_word_ids): prev_word_idx = int(prev_word_ids[old_slot]) if prev_word_idx >= 0: votes[prev_word_idx] += 1 strict = [key for key, count in votes.items() if count > len(positions) / 2.0] decision = { 'candidate_prev_word_id': None, 'selected_prev_word_id': None, 'predecessor_word_ids': sorted(int(key) for key in votes), 'vote_counts': {str(key): int(value) for key, value in sorted(votes.items())}, 'mapping_status': 'unmatched_new' if not votes else 'ambiguous_no_majority', 'split_competitor_count': 0, 'reliable': not votes, 'text_equal': False, 'time_iou': 0.0, } if len(strict) == 1: prev_word_idx = int(strict[0]) prev_meta = prev_word_states.get(prev_word_idx, {}).get('metadata', {}) cur_meta = _commit_word_metadata(word_timestamps, word_idx) decision.update({ 'candidate_prev_word_id': prev_word_idx, 'mapping_status': 'candidate', 'text_equal': bool(prev_meta.get('text') and cur_meta['text'] and prev_meta['text'] == cur_meta['text']), 'time_iou': float(_commit_time_iou(prev_meta, cur_meta)), }) competitors_by_prev.setdefault(prev_word_idx, []).append((word_idx, decision)) decisions[word_idx] = decision # A strict-majority edge is unique for each current word. Selecting the # highest-weight edge for each predecessor is therefore the global maximum # weight one-to-one matching for this candidate graph. for prev_word_idx, competitors in sorted(competitors_by_prev.items()): for _, decision in competitors: decision['split_competitor_count'] = len(competitors) ranked = sorted( competitors, key=lambda item: ( -int(item[1]['vote_counts'][str(prev_word_idx)]), -round(float(item[1]['time_iou']), 12), -int(item[1]['text_equal']), int(item[0]), ), ) winner_key = ( int(ranked[0][1]['vote_counts'][str(prev_word_idx)]), round(float(ranked[0][1]['time_iou']), 12), int(ranked[0][1]['text_equal']), ) tied = [ item for item in ranked if ( int(item[1]['vote_counts'][str(prev_word_idx)]), round(float(item[1]['time_iou']), 12), int(item[1]['text_equal']), ) == winner_key ] if len(tied) == 1: _, winner = ranked[0] winner['selected_prev_word_id'] = int(prev_word_idx) winner['mapping_status'] = 'accepted' winner['reliable'] = True else: for _, decision in tied: decision['mapping_status'] = 'ambiguous_competition_tie' decision['reliable'] = False for _, loser in ranked[1:]: if loser['mapping_status'] == 'candidate': loser['mapping_status'] = 'split_new_successor' loser['reliable'] = True return current_positions, decisions def build_stateful_commit_masks( prev_state, pcn, pcn_word_id, word_timestamps, commit_time, audio_end, is_final, current_to_prev_override=None, commit_rule='legacy', ): if commit_rule not in {'legacy', 'merge_aware'}: raise ValueError(f"commit_rule must be 'legacy' or 'merge_aware', got {commit_rule!r}") word_timestamps = _top_word_timestamps(word_timestamps) candidate = cumulative_candidate_mask_from_asr(pcn_word_id, word_timestamps, commit_time, audio_end, is_final) if current_to_prev_override is None: prev_to_current, mapped_old_slots, dropped_or_revised = align_previous_commits(prev_state, pcn['top_phone_ids']) current_to_prev = np.zeros((len(pcn['top_phone_ids']),), dtype=np.int32) - 1 for prev_idx, cur_idx in prev_to_current.items(): if 0 <= cur_idx < current_to_prev.shape[0]: current_to_prev[cur_idx] = int(prev_idx) else: current_to_prev = np.asarray(current_to_prev_override, dtype=np.int32).copy() prev_to_current = { int(prev_idx): int(cur_idx) for cur_idx, prev_idx in enumerate(current_to_prev.tolist()) if int(prev_idx) >= 0 } prev_committed = set(int(value) for value in (prev_state or {}).get('committed_slots', [])) mapped_old_slots = sorted(cur for prev, cur in prev_to_current.items() if prev in prev_committed) dropped_or_revised = sorted(prev for prev in prev_committed if prev not in prev_to_current) mapped_old_set = set(mapped_old_slots) cumulative = candidate.copy() for slot_idx in mapped_old_set: if 0 <= slot_idx < cumulative.shape[0]: cumulative[slot_idx] = 1.0 if commit_rule == 'legacy': new_mask = np.zeros_like(cumulative) for word_idx in sorted(set(int(value) for value in pcn_word_id.tolist() if int(value) >= 0)): positions = np.flatnonzero(pcn_word_id == word_idx) if positions.size == 0 or not bool(np.all(candidate[positions] > 0)): continue if any(int(position) in mapped_old_set for position in positions.tolist()): continue new_mask[positions] = 1.0 committed_slots = sorted(np.flatnonzero(cumulative > 0).astype(int).tolist()) new_slots = sorted(np.flatnonzero(new_mask > 0).astype(int).tolist()) diagnostics = { 'mapped_old_slots': mapped_old_slots, 'new_slots': new_slots, 'dropped_or_revised_slots': dropped_or_revised, 'state_event_semantics': 'legacy_slot_any_v1', } next_state = { 'top_phone_ids': list(pcn['top_phone_ids']), 'committed_slots': committed_slots, } return cumulative.astype(np.float32), new_mask.astype(np.float32), current_to_prev, diagnostics, next_state current_positions, lineage = _word_level_lineage( prev_state, np.asarray(pcn_word_id, dtype=np.int32), current_to_prev, word_timestamps, ) prev_word_states = (prev_state or {}).get('word_states', {}) next_logical_id = int((prev_state or {}).get('next_logical_word_id', 0)) next_event_id = int((prev_state or {}).get('next_state_event_id', 0)) next_word_states = {} word_event_rows = [] new_mask = np.zeros_like(cumulative) for word_idx, positions in sorted(current_positions.items()): decision = lineage[word_idx] selected_prev = decision['selected_prev_word_id'] predecessor_ids = decision['predecessor_word_ids'] predecessor_states = [prev_word_states.get(value, {}) for value in predecessor_ids] predecessor_committed = [bool(value.get('committed', False)) for value in predecessor_states] completed = bool(np.all(candidate[positions] > 0)) split = int(decision['split_competitor_count']) > 1 merge = len(predecessor_ids) > 1 inherited_state = prev_word_states.get(selected_prev, {}) if selected_prev is not None else {} has_old_committed_slot = any(int(position) in mapped_old_set for position in positions.tolist()) prev_word_ids_array = np.asarray((prev_state or {}).get('pcn_word_id', []), dtype=np.int32) prev_committed_slots = set(int(value) for value in (prev_state or {}).get('committed_slots', [])) unassigned_committed_predecessors = 0 for position in positions.tolist(): old_slot = int(current_to_prev[position]) if position < len(current_to_prev) else -1 if old_slot not in prev_committed_slots: continue if old_slot >= len(prev_word_ids_array) or int(prev_word_ids_array[old_slot]) < 0: unassigned_committed_predecessors += 1 if unassigned_committed_predecessors: decision['mapping_reliable'] = False decision['reliable'] = False decision['mapping_status'] = 'lineage_broken_unassigned_committed_slot' legacy_emit = completed and not has_old_committed_slot mixed_merge = bool( merge and predecessor_committed and any(predecessor_committed) and not all(predecessor_committed) ) partial_word_completed = bool( completed and selected_prev is not None and not inherited_state.get('committed', False) and not merge and bool(decision['reliable']) ) # A3-Lite is intentionally minimal: reproduce the legacy mask exactly, # then add one event for a completed mixed committed/partial merge or # for a reliably inherited partial word that has now become complete. emit_event = bool(legacy_emit or (completed and mixed_merge) or partial_word_completed) committed = bool(completed or inherited_state.get('committed', False)) if selected_prev is not None and not (completed and mixed_merge): logical_id = int(inherited_state['logical_word_id']) else: logical_id = next_logical_id next_logical_id += 1 if completed and mixed_merge: event_reason = 'mixed_committed_partial_merge_completed' elif mixed_merge: event_reason = 'mixed_committed_partial_merge_partial' elif split: event_reason = 'split_completed' if completed else 'split_partial' elif merge: if completed: event_reason = 'merge_completed' if emit_event else 'merge_completed_legacy_suppressed' else: event_reason = 'merge_partial' elif unassigned_committed_predecessors and completed and not emit_event: event_reason = 'lineage_broken_completed_legacy_suppressed' elif selected_prev is not None and inherited_state.get('committed', False): event_reason = 'persistent_committed' elif selected_prev is not None: event_reason = 'partial_became_completed' if completed else 'partial_persistent' elif predecessor_ids: event_reason = 'ambiguous_completed' if completed else 'ambiguous_partial' else: event_reason = 'newly_completed' if completed else 'new_partial' state_event_id = -1 if emit_event: new_mask[positions] = 1.0 state_event_id = next_event_id next_event_id += 1 metadata = _commit_word_metadata(word_timestamps, word_idx) next_word_states[int(word_idx)] = { 'logical_word_id': int(logical_id), 'committed': bool(committed), 'state_event_id': int(state_event_id if state_event_id >= 0 else inherited_state.get('state_event_id', -1)), 'metadata': metadata, } word_event_rows.append({ 'current_word_id': int(word_idx), 'logical_word_id': int(logical_id), 'state_event_id': int(state_event_id), 'completed': bool(completed), 'emit_state_event': bool(emit_event), 'event_reason': event_reason, 'mapping_status': decision['mapping_status'], 'mapping_reliable': bool(decision['reliable']), 'selected_prev_word_id': selected_prev, 'predecessor_word_ids': predecessor_ids, 'predecessor_committed': predecessor_committed, 'vote_counts': decision['vote_counts'], 'split_competitor_count': int(decision['split_competitor_count']), 'text': metadata['text'], 'text_equal': bool(decision['text_equal']), 'time_iou': float(decision['time_iou']), 'unassigned_committed_predecessor_slots': int(unassigned_committed_predecessors), }) committed_slots = sorted(np.flatnonzero(cumulative > 0).astype(int).tolist()) new_slots = sorted(np.flatnonzero(new_mask > 0).astype(int).tolist()) diagnostics = { 'mapped_old_slots': mapped_old_slots, 'new_slots': new_slots, 'dropped_or_revised_slots': dropped_or_revised, 'word_level_state_events': word_event_rows, 'state_event_semantics': 'word_predecessor_one_to_one_v1', } next_state = { 'top_phone_ids': list(pcn['top_phone_ids']), 'pcn_word_id': np.asarray(pcn_word_id, dtype=np.int32).tolist(), 'committed_slots': committed_slots, 'word_states': next_word_states, 'next_logical_word_id': int(next_logical_id), 'next_state_event_id': int(next_event_id), } return cumulative.astype(np.float32), new_mask.astype(np.float32), current_to_prev, diagnostics, next_state def acoustic_distribution_for_frames(frame_probs, phone_to_frame_id, phn_dict, phone_dim): out = np.zeros((phone_dim,), dtype=np.float32) if frame_probs.size == 0: return out mean_probs = frame_probs.mean(axis=0) for phone, phone_idx in phn_dict.items(): frame_idx = phone_to_frame_id.get(phone) if frame_idx is not None and frame_idx < mean_probs.shape[0]: out[int(phone_idx)] = float(mean_probs[frame_idx]) if out.sum() > 0: out /= out.sum() return out def build_acoustic_evidence(item, cn_post, top_phone_ids, phone_to_frame_id, phn_dict, audio_end): phone_dim = cn_post.shape[-1] acoustic_post = np.zeros_like(cn_post, dtype=np.float32) acoustic_stats = np.zeros((cn_post.shape[0], 4), dtype=np.float32) kept_indices, kept_probs = select_visible_frames(item['probs'], item['keep_mask'], audio_end, item['frame_step']) alignable = [ (slot_idx, phone_idx) for slot_idx, phone_idx in enumerate(top_phone_ids) if phone_idx is not None ] if kept_probs.shape[0] < len(alignable) or not alignable: for slot_idx in range(cn_post.shape[0]): acoustic_stats[slot_idx, 3] = js_divergence(cn_post[slot_idx], acoustic_post[slot_idx]) return acoustic_post, acoustic_stats, 0 id_to_phone = {idx: phone for phone, idx in phn_dict.items()} frame_phone_ids = [phone_to_frame_id[id_to_phone[phone_idx]] for _, phone_idx in alignable] try: path = monotonic_align(-np.log(np.clip(kept_probs, EPS, None)), frame_phone_ids) except Exception: for slot_idx in range(cn_post.shape[0]): acoustic_stats[slot_idx, 3] = js_divergence(cn_post[slot_idx], acoustic_post[slot_idx]) return acoustic_post, acoustic_stats, 0 for local_idx, (slot_idx, _) in enumerate(alignable): tok_frames = kept_indices[path == local_idx] if tok_frames.size == 0: acoustic_stats[slot_idx, 3] = js_divergence(cn_post[slot_idx], acoustic_post[slot_idx]) continue dist = acoustic_distribution_for_frames( item['probs'][tok_frames], phone_to_frame_id, phn_dict, phone_dim, ) acoustic_post[slot_idx] = dist non_eps = dist[:-1] _, margin = top_margin(non_eps) duration = float(tok_frames.size * item['frame_step']) acoustic_stats[slot_idx] = np.array( [ entropy_np(dist), margin, duration, js_divergence(cn_post[slot_idx], dist), ], dtype=np.float32, ) return acoustic_post, acoustic_stats, int(kept_probs.shape[0]) def compute_prosody(audio, sample_rate, audio_end, word_count, phone_count): audio = np.asarray(audio, dtype=np.float32) duration = float(max(audio_end, EPS)) if audio.size == 0: return np.zeros((14,), dtype=np.float32) hop_length = max(int(0.01 * sample_rate), 1) frame_length = max(int(0.025 * sample_rate), hop_length) rms = librosa.feature.rms(y=audio, frame_length=frame_length, hop_length=hop_length)[0] log_energy = np.log(np.clip(rms, EPS, None)) energy_threshold = np.percentile(rms, 20) if rms.size else 0.0 silence = rms <= max(energy_threshold, EPS) silence_ratio = float(np.mean(silence)) if silence.size else 0.0 pause_count, longest_pause = count_pauses(silence, hop_length / sample_rate) try: f0, voiced_flag, voiced_prob = librosa.pyin( audio, fmin=50, fmax=500, sr=sample_rate, frame_length=frame_length, hop_length=hop_length, ) voiced = np.isfinite(f0) f0_values = f0[voiced] voiced_prob_mean = float(np.nanmean(voiced_prob)) if voiced_prob is not None else float(np.mean(voiced)) except Exception: f0_values = np.array([], dtype=np.float32) voiced = np.zeros_like(rms, dtype=bool) voiced_prob_mean = 0.0 f0_mean = float(np.mean(f0_values)) if f0_values.size else 0.0 f0_std = float(np.std(f0_values)) if f0_values.size else 0.0 f0_slope = slope(f0_values) energy_mean = float(np.mean(log_energy)) if log_energy.size else 0.0 energy_std = float(np.std(log_energy)) if log_energy.size else 0.0 energy_slope = slope(log_energy) voiced_duration = max(float(np.mean(voiced)) * duration, EPS) if voiced.size else EPS return np.array( [ duration, f0_mean, f0_std, f0_slope, voiced_prob_mean, energy_mean, energy_std, energy_slope, silence_ratio, float(pause_count), float(longest_pause), float(word_count) / duration, float(phone_count) / duration, float(phone_count) / voiced_duration, ], dtype=np.float32, ) def slot_prosody_feature_names(): return [ 'slot_duration', 'slot_log_energy_mean', 'slot_log_energy_std', 'slot_log_energy_max', 'slot_f0_mean', 'slot_f0_std', 'slot_f0_max', 'slot_voiced_ratio', 'slot_position_in_word', 'word_phone_count', 'energy_relative_to_word_mean', 'duration_relative_to_word_mean', 'f0_relative_to_word_mean', 'is_vowel', 'lexical_stress_0', 'lexical_stress_1', 'lexical_stress_2', ] def compute_slot_prosody(audio, sample_rate, slot_times, pcn_word_id, top_phone_ids, phn_dict): feature_names = slot_prosody_feature_names() slot_count = len(top_phone_ids) out = np.zeros((slot_count, len(feature_names)), dtype=np.float32) if slot_count <= 0: return out audio = np.asarray(audio, dtype=np.float32) id_to_phone = {int(idx): phone for phone, idx in phn_dict.items()} hop_length = max(int(0.01 * sample_rate), 1) frame_length = max(int(0.025 * sample_rate), hop_length) if audio.size: rms = librosa.feature.rms(y=audio, frame_length=frame_length, hop_length=hop_length)[0] log_energy = np.log(np.clip(rms, EPS, None)) frame_times = librosa.frames_to_time(np.arange(rms.shape[0]), sr=sample_rate, hop_length=hop_length) else: log_energy = np.zeros((0,), dtype=np.float32) frame_times = np.zeros((0,), dtype=np.float32) try: f0, _, _ = librosa.pyin( audio, fmin=50, fmax=500, sr=sample_rate, frame_length=frame_length, hop_length=hop_length, ) except Exception: f0 = np.zeros_like(log_energy, dtype=np.float32) * np.nan if f0 is None: f0 = np.zeros_like(log_energy, dtype=np.float32) * np.nan f0 = np.asarray(f0, dtype=np.float32) common_frames = min(frame_times.shape[0], log_energy.shape[0], f0.shape[0]) frame_times = frame_times[:common_frames] log_energy = log_energy[:common_frames] f0 = f0[:common_frames] durations = np.zeros((slot_count,), dtype=np.float32) energy_means = np.zeros((slot_count,), dtype=np.float32) f0_means = np.zeros((slot_count,), dtype=np.float32) pcn_word_id = np.asarray(pcn_word_id, dtype=np.int32) for slot_idx in range(slot_count): if slot_idx < len(slot_times) and slot_times[slot_idx] is not None: start, end = slot_times[slot_idx] start = float(start) end = float(max(end, start)) else: start = 0.0 end = 0.0 durations[slot_idx] = max(end - start, 0.0) frame_mask = (frame_times >= start) & (frame_times <= end) if frame_times.size else np.zeros((0,), dtype=bool) energy_values = log_energy[frame_mask] f0_values = f0[frame_mask] f0_values = f0_values[np.isfinite(f0_values)] if energy_values.size: out[slot_idx, 1] = float(np.mean(energy_values)) out[slot_idx, 2] = float(np.std(energy_values)) out[slot_idx, 3] = float(np.max(energy_values)) energy_means[slot_idx] = out[slot_idx, 1] if f0_values.size: out[slot_idx, 4] = float(np.mean(f0_values)) out[slot_idx, 5] = float(np.std(f0_values)) out[slot_idx, 6] = float(np.max(f0_values)) out[slot_idx, 7] = float(f0_values.size / max(frame_mask.sum(), 1)) f0_means[slot_idx] = out[slot_idx, 4] out[slot_idx, 0] = durations[slot_idx] word_idx = int(pcn_word_id[slot_idx]) if slot_idx < pcn_word_id.shape[0] else -1 same_word = np.flatnonzero(pcn_word_id == word_idx) if word_idx >= 0 else np.asarray([slot_idx], dtype=np.int64) if same_word.size: position = np.flatnonzero(same_word == slot_idx) out[slot_idx, 8] = float(position[0] / max(same_word.size - 1, 1)) if position.size else 0.0 out[slot_idx, 9] = float(same_word.size) phone_idx = top_phone_ids[slot_idx] if slot_idx < len(top_phone_ids) else None phone = id_to_phone.get(int(phone_idx), '') if phone_idx is not None else '' out[slot_idx, 13] = 1.0 if phone[:1].upper() in {'A', 'E', 'I', 'O', 'U'} else 0.0 for slot_idx in range(slot_count): word_idx = int(pcn_word_id[slot_idx]) if slot_idx < pcn_word_id.shape[0] else -1 same_word = np.flatnonzero(pcn_word_id == word_idx) if word_idx >= 0 else np.asarray([slot_idx], dtype=np.int64) if same_word.size: word_energy = float(np.mean(energy_means[same_word])) word_duration = float(np.mean(durations[same_word])) word_f0 = float(np.mean(f0_means[same_word])) out[slot_idx, 10] = out[slot_idx, 1] - word_energy out[slot_idx, 11] = durations[slot_idx] / max(word_duration, EPS) out[slot_idx, 12] = out[slot_idx, 4] - word_f0 return out.astype(np.float32) def slope(values): values = np.asarray(values, dtype=np.float32) if values.size < 2: return 0.0 x = np.arange(values.size, dtype=np.float32) x = x - x.mean() y = values - values.mean() denom = float(np.sum(x * x)) if denom <= EPS: return 0.0 return float(np.sum(x * y) / denom) def count_pauses(silence_mask, frame_step): pause_count = 0 longest = 0.0 cur = 0 for value in silence_mask.tolist(): if value: cur += 1 elif cur > 0: pause_count += 1 longest = max(longest, cur * frame_step) cur = 0 if cur > 0: pause_count += 1 longest = max(longest, cur * frame_step) return pause_count, longest