mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
Grapette.LandLesereingrape
22b42827a0
A blank or whitespace-only line in a file list crashes prepare_data_iterator: the id/data branch indexes lines[0] on an empty split (IndexError), and the .jsonl branch feeds an empty string to json.loads (JSONDecodeError). Skip blank lines so a trailing or separator line no longer aborts the whole batch. Co-authored-by: Lesereingrape <[email protected]>
1522 lines
64 KiB
Python
1522 lines
64 KiB
Python
#!/usr/bin/env python3
|
||
# -*- encoding: utf-8 -*-
|
||
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
|
||
# MIT License (https://opensource.org/licenses/MIT)
|
||
|
||
import json
|
||
import time
|
||
import copy
|
||
import torch
|
||
import random
|
||
import re
|
||
import string
|
||
import logging
|
||
import os.path
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
|
||
from omegaconf import DictConfig, ListConfig
|
||
from funasr.utils.misc import deep_update
|
||
from funasr.register import tables
|
||
from funasr.utils.load_utils import load_bytes
|
||
from funasr.download.file import download_from_url
|
||
from funasr.utils.timestamp_tools import timestamp_sentence
|
||
from funasr.utils.timestamp_tools import timestamp_sentence_en
|
||
from funasr.download.download_model_from_hub import download_model
|
||
from funasr.utils.vad_utils import slice_padding_audio_samples
|
||
from funasr.utils.vad_utils import merge_vad
|
||
from funasr.utils.load_utils import load_audio_text_image_video
|
||
from funasr.train_utils.set_all_random_seed import set_all_random_seed
|
||
from funasr.train_utils.load_pretrained_model import load_pretrained_model
|
||
from funasr.utils import export_utils
|
||
from funasr.utils.postprocess_hotwords import apply_postprocess_hotwords_to_results
|
||
from funasr.utils import misc
|
||
|
||
|
||
def is_npu_available():
|
||
"""检查NPU是否可用。"""
|
||
try:
|
||
import torch_npu
|
||
|
||
return torch_npu.npu.is_available()
|
||
except ImportError:
|
||
return False
|
||
|
||
|
||
def _resolve_ncpu(config, fallback=4):
|
||
"""Return a positive integer representing CPU threads from config."""
|
||
value = config.get("ncpu", fallback)
|
||
try:
|
||
value = int(value)
|
||
except (TypeError, ValueError):
|
||
value = fallback
|
||
return max(value, 1)
|
||
|
||
|
||
# The active BLAS limit, if any. Held for the life of the process rather than
|
||
# scoped to a call -- see `_limit_blas_threads` for why.
|
||
_blas_thread_limiter = None
|
||
_blas_thread_limit = None
|
||
|
||
|
||
def _current_blas_threads():
|
||
"""Thread count currently configured for BLAS, or None if unreadable."""
|
||
try:
|
||
import threadpoolctl
|
||
except ImportError:
|
||
return None
|
||
counts = {
|
||
pool["num_threads"]
|
||
for pool in threadpoolctl.threadpool_info()
|
||
if pool["user_api"] == "blas"
|
||
}
|
||
return counts.pop() if len(counts) == 1 else None
|
||
|
||
|
||
def _limit_blas_threads(ncpu):
|
||
"""Apply `ncpu` to BLAS as well as to torch.
|
||
|
||
`ncpu` documents itself as the thread count for "CPU 内部操作并行性", but
|
||
it only ever reached `torch.set_num_threads`. The speaker-clustering path
|
||
does not run on torch at all: `scipy.linalg.eigh` on the affinity
|
||
Laplacian goes to BLAS, whose own default is one thread per core. So on a
|
||
many-core host the CPU is saturated by a library `ncpu` never touched.
|
||
Measured on a 64-core host with a 357-row Laplacian:
|
||
|
||
BLAS threads wall CPU
|
||
64 1.20s 68.9 core-seconds
|
||
4 0.05s 2.8 core-seconds
|
||
1 0.03s 1.8 core-seconds
|
||
|
||
At a few hundred rows the parallel driver spends its time synchronising
|
||
rather than computing, so bounding it is a latency win as well as a CPU
|
||
one.
|
||
|
||
The update follows `torch.set_num_threads`: the newest `ncpu` wins, and no
|
||
set is issued when BLAS already sits at that value. Both settings are
|
||
process-wide, so constructing a second `AutoModel` applies its `ncpu` to
|
||
any model already built in this process -- including ones without speaker
|
||
clustering. That is the same contract `torch.set_num_threads` already
|
||
has; the difference is only that `ncpu` used to leave BLAS untouched.
|
||
|
||
The limit is not released once applied. It is process-global, so a scoped
|
||
enter/exit pair per request could be interleaved by overlapping requests
|
||
and leave the setting unbalanced, and FunASR reaches BLAS only through
|
||
small one-shot operations. Doing nothing when threadpoolctl is missing
|
||
leaves BLAS at its own default rather than failing model construction.
|
||
"""
|
||
global _blas_thread_limiter, _blas_thread_limit
|
||
|
||
current = _current_blas_threads()
|
||
if current is not None and current == ncpu:
|
||
return
|
||
|
||
# `threadpool_limits` records the OpenMP pools as well as the BLAS ones
|
||
# even when entered with `user_api="blas"`, and its `__exit__` restores
|
||
# every pool it captured. Dropping the previous limiter would therefore
|
||
# roll torch's thread count back to whatever it was when that limiter was
|
||
# created -- undoing the `torch.set_num_threads(ncpu)` that `build_model`
|
||
# has already performed. Save and restore it across the swap so replacing
|
||
# the BLAS limit cannot touch an OpenMP setting it does not own.
|
||
torch_threads = torch.get_num_threads()
|
||
|
||
if _blas_thread_limiter is not None:
|
||
_blas_thread_limiter.__exit__(None, None, None)
|
||
_blas_thread_limiter = None
|
||
|
||
try:
|
||
import threadpoolctl
|
||
except ImportError:
|
||
logging.info(
|
||
"threadpoolctl is not installed, so `ncpu` cannot be applied to "
|
||
"BLAS; speaker clustering may use one thread per core."
|
||
)
|
||
return
|
||
|
||
_blas_thread_limiter = threadpoolctl.threadpool_limits(
|
||
limits=ncpu, user_api="blas"
|
||
)
|
||
_blas_thread_limiter.__enter__()
|
||
_blas_thread_limit = ncpu
|
||
|
||
if torch.get_num_threads() != torch_threads:
|
||
torch.set_num_threads(torch_threads)
|
||
|
||
# BLAS pools are only visible to threadpoolctl once a library backed by
|
||
# them has been loaded. funasr imports such a library at package import,
|
||
# so they are present by the time a model is built -- but confirm rather
|
||
# than assume, because a dropped limit here would silently restore the
|
||
# CPU saturation this is meant to prevent.
|
||
applied = _current_blas_threads()
|
||
if applied is not None and applied != ncpu:
|
||
logging.warning(
|
||
"ncpu=%s could not be applied to BLAS (it reports %s); "
|
||
"speaker clustering may use more threads than requested.",
|
||
ncpu,
|
||
applied,
|
||
)
|
||
|
||
|
||
def _join_vad_texts(texts):
|
||
"""Remove rich tags and join VAD text without adding spaces between Chinese chunks."""
|
||
cleaned = [re.sub(r"<\|[^|]*\|>", "", text).strip() for text in texts]
|
||
cleaned = [text for text in cleaned if text]
|
||
if not cleaned:
|
||
return ""
|
||
joined = cleaned[0]
|
||
for text in cleaned[1:]:
|
||
separator = ""
|
||
if not ("\u3400" <= joined[-1] <= "\u9fff" and "\u3400" <= text[0] <= "\u9fff"):
|
||
separator = " "
|
||
joined += separator + text
|
||
return joined
|
||
|
||
|
||
def _vad_segment_sentences(restored_data, vadsegments):
|
||
"""Build readable sentence records directly from VAD-aligned ASR chunks."""
|
||
sentences = []
|
||
for result, vadsegment in zip(restored_data, vadsegments):
|
||
text = re.sub(r"<\|[^|]*\|>", "", str(result.get("text", ""))).strip()
|
||
if not text:
|
||
continue
|
||
|
||
timestamps = []
|
||
raw_timestamps = result.get("timestamp")
|
||
if raw_timestamps is None:
|
||
raw_timestamps = result.get("timestamps", [])
|
||
for item in raw_timestamps or []:
|
||
if isinstance(item, dict):
|
||
start = item.get("start_time")
|
||
end = item.get("end_time")
|
||
if start is None or end is None:
|
||
continue
|
||
timestamps.append([int(float(start) * 1000), int(float(end) * 1000)])
|
||
elif isinstance(item, (list, tuple)) and len(item) >= 2:
|
||
timestamps.append([int(item[0]), int(item[1])])
|
||
|
||
start = timestamps[0][0] if timestamps else vadsegment[0]
|
||
end = timestamps[-1][1] if timestamps else vadsegment[1]
|
||
sentences.append(
|
||
{
|
||
"start": start,
|
||
"end": end,
|
||
"text": text,
|
||
"sentence": text,
|
||
"timestamp": timestamps,
|
||
}
|
||
)
|
||
return sentences
|
||
|
||
|
||
def _get_punc_tokens(text, punc_array, punc_model):
|
||
"""Return the surface tokens represented by a CT-Transformer punctuation array."""
|
||
try:
|
||
from funasr.models.ct_transformer.utils import split_words
|
||
|
||
tokens = split_words(
|
||
text,
|
||
jieba_usr_dict=getattr(punc_model, "jieba_usr_dict", None),
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
expanded_tokens = []
|
||
for token in tokens:
|
||
if token and "\u0e00" <= token[0] <= "\u9fa5" and len(token) > 1:
|
||
expanded_tokens.extend(token)
|
||
else:
|
||
expanded_tokens.append(token)
|
||
try:
|
||
punc_length = len(punc_array)
|
||
except TypeError:
|
||
return None
|
||
if len(expanded_tokens) != punc_length:
|
||
return None
|
||
return expanded_tokens
|
||
|
||
|
||
def _punctuate_surface_text(text, punc_array, punc_model):
|
||
"""Insert predicted punctuation without changing the ASR surface text."""
|
||
tokens = _get_punc_tokens(text, punc_array, punc_model)
|
||
if tokens is None:
|
||
return None
|
||
|
||
spans = _surface_token_spans(text, tokens)
|
||
if spans is None:
|
||
return None
|
||
|
||
parts = []
|
||
cursor = 0
|
||
for token, punc_id, (_, end) in zip(tokens, punc_array, spans):
|
||
parts.append(text[cursor:end])
|
||
parts.append(_punc_symbol(punc_id, token, punc_model))
|
||
cursor = end
|
||
parts.append(text[cursor:])
|
||
return "".join(parts)
|
||
|
||
|
||
def _punc_symbol(punc_id, token, punc_model):
|
||
"""Return the punctuation character represented by a model punctuation ID."""
|
||
punc_list = getattr(punc_model, "punc_list", None)
|
||
fallback_punc = {1: "", 2: "，", 3: "。", 4: "？", 5: "、"}
|
||
punc_id = int(punc_id)
|
||
try:
|
||
punctuation = punc_list[punc_id]
|
||
except (IndexError, TypeError):
|
||
punctuation = fallback_punc.get(punc_id, "")
|
||
if punctuation == "_":
|
||
punctuation = ""
|
||
if punctuation and token[0].isascii():
|
||
punctuation = {"，": ",", "。": ".", "？": "?", "、": ","}.get(
|
||
punctuation, punctuation
|
||
)
|
||
return punctuation
|
||
|
||
|
||
def _surface_token_spans(text, tokens):
|
||
"""Map punctuation tokens back to exact spans in the original surface text."""
|
||
spans = []
|
||
cursor = 0
|
||
for token in tokens:
|
||
while cursor < len(text) and text[cursor].isspace():
|
||
cursor += 1
|
||
surface_token = text[cursor : cursor + len(token)]
|
||
if surface_token.casefold() != token.casefold():
|
||
return None
|
||
spans.append((cursor, cursor + len(token)))
|
||
cursor += len(token)
|
||
if text[cursor:].strip():
|
||
return None
|
||
return spans
|
||
|
||
|
||
def _merge_timestamp_units(text, words, timestamps, punc_array, punc_model):
|
||
"""Merge timestamp/BPE units to the punctuation model's surface tokens."""
|
||
expanded_tokens = _get_punc_tokens(text, punc_array, punc_model)
|
||
if expanded_tokens is None:
|
||
return None
|
||
|
||
def normalize(value):
|
||
return "".join(value.split()).replace("▁", "").casefold()
|
||
|
||
if len(words) != len(timestamps):
|
||
return None
|
||
|
||
aligned_text = ""
|
||
character_timestamps = []
|
||
for word, timestamp in zip(words, timestamps):
|
||
word_text = normalize(word)
|
||
if (
|
||
not word_text
|
||
or not isinstance(timestamp, (list, tuple))
|
||
or len(timestamp) < 2
|
||
or timestamp[1] < timestamp[0]
|
||
):
|
||
return None
|
||
start, end = timestamp[:2]
|
||
duration = end - start
|
||
aligned_text += word_text
|
||
for index in range(len(word_text)):
|
||
character_timestamps.append(
|
||
[
|
||
start + duration * index // len(word_text),
|
||
start + duration * (index + 1) // len(word_text),
|
||
]
|
||
)
|
||
|
||
merged_timestamps = []
|
||
character_index = 0
|
||
for token in expanded_tokens:
|
||
token_text = normalize(token)
|
||
token_end = character_index + len(token_text)
|
||
if (
|
||
not token_text
|
||
or aligned_text[character_index:token_end] != token_text
|
||
or token_end > len(character_timestamps)
|
||
):
|
||
return None
|
||
merged_timestamps.append(
|
||
[
|
||
character_timestamps[character_index][0],
|
||
character_timestamps[token_end - 1][1],
|
||
]
|
||
)
|
||
character_index = token_end
|
||
|
||
if character_index != len(character_timestamps):
|
||
return None
|
||
return " ".join(expanded_tokens), merged_timestamps
|
||
|
||
|
||
def _timestamp_sentences_from_surface(
|
||
text, timestamps, punc_array, punc_model, return_raw_text=False
|
||
):
|
||
"""Build sentence timestamps while preserving the exact ASR surface text."""
|
||
tokens = _get_punc_tokens(text, punc_array, punc_model)
|
||
if tokens is None or len(tokens) != len(timestamps):
|
||
return None
|
||
spans = _surface_token_spans(text, tokens)
|
||
if spans is None:
|
||
return None
|
||
|
||
sentences = []
|
||
sentence_start = 0
|
||
for index, (token, punc_id) in enumerate(zip(tokens, punc_array)):
|
||
punctuation = _punc_symbol(punc_id, token, punc_model)
|
||
if not punctuation:
|
||
continue
|
||
raw_sentence = text[spans[sentence_start][0] : spans[index][1]].strip()
|
||
sentence = {
|
||
"text": raw_sentence + punctuation,
|
||
"start": timestamps[sentence_start][0],
|
||
"end": timestamps[index][1],
|
||
"timestamp": timestamps[sentence_start : index + 1],
|
||
}
|
||
if return_raw_text:
|
||
sentence["raw_text"] = raw_sentence
|
||
sentences.append(sentence)
|
||
sentence_start = index + 1
|
||
|
||
if sentence_start < len(tokens):
|
||
raw_sentence = text[spans[sentence_start][0] : spans[-1][1]].strip()
|
||
sentence = {
|
||
"text": raw_sentence,
|
||
"start": timestamps[sentence_start][0],
|
||
"end": timestamps[-1][1],
|
||
"timestamp": timestamps[sentence_start:],
|
||
}
|
||
if return_raw_text:
|
||
sentence["raw_text"] = raw_sentence
|
||
sentences.append(sentence)
|
||
return sentences
|
||
|
||
|
||
def _get_import_errors():
|
||
"""Internal: get import errors."""
|
||
try:
|
||
import funasr
|
||
except Exception:
|
||
return {}
|
||
get_import_errors = getattr(funasr, "get_import_errors", None)
|
||
if get_import_errors is not None:
|
||
return get_import_errors()
|
||
return dict(getattr(funasr, "_IMPORT_ERRORS", {}))
|
||
|
||
|
||
def _format_unregistered_component_error(component_type, component_name, registry):
|
||
"""Internal: format unregistered component error.
|
||
|
||
Args:
|
||
component_type: TODO.
|
||
component_name: TODO.
|
||
registry: TODO.
|
||
"""
|
||
registered = sorted(registry.keys())
|
||
preview = ", ".join(registered[:80])
|
||
if len(registered) > 80:
|
||
preview += f", ... ({len(registered)} total)"
|
||
if not preview:
|
||
preview = "(none)"
|
||
|
||
import_errors = _get_import_errors()
|
||
if import_errors:
|
||
lines = [
|
||
f" - {name}: {error}"
|
||
for name, error in sorted(import_errors.items())[:50]
|
||
]
|
||
remaining = len(import_errors) - len(lines)
|
||
if remaining > 0:
|
||
lines.append(f" ... {remaining} more import failures hidden")
|
||
import_error_text = "\n".join(lines)
|
||
else:
|
||
import_error_text = " (none recorded)"
|
||
|
||
return (
|
||
f"{component_type} '{component_name}' is not registered.\n"
|
||
f"Registered {component_type} keys ({len(registered)}): {preview}\n"
|
||
"Some modules may have failed to import during auto-registration. "
|
||
"Set FUNASR_IMPORT_DEBUG=1 to print failures during import, or "
|
||
"FUNASR_STRICT_IMPORT=1 to fail fast.\n"
|
||
f"Recorded import failures:\n{import_error_text}"
|
||
)
|
||
|
||
|
||
try:
|
||
from funasr.models.campplus.utils import sv_chunk, postprocess, distribute_spk
|
||
from funasr.models.campplus.cluster_backend import ClusterBackend
|
||
except:
|
||
pass
|
||
|
||
|
||
def prepare_data_iterator(
|
||
data_in, input_len=None, data_type=None, key=None, input_format="auto"
|
||
):
|
||
"""Prepare inputs, optionally interpreting audio bytes as explicit PCM16."""
|
||
data_list = []
|
||
key_list = []
|
||
filelist = [".scp", ".txt", ".json", ".jsonl", ".text"]
|
||
|
||
chars = string.ascii_letters + string.digits
|
||
if isinstance(data_in, str):
|
||
if data_in.startswith("http://") or data_in.startswith("https://"): # url
|
||
data_in = download_from_url(data_in)
|
||
|
||
if isinstance(data_in, str) and os.path.exists(
|
||
data_in
|
||
): # wav_path; filelist: wav.scp, file.jsonl;text.txt;
|
||
_, file_extension = os.path.splitext(data_in)
|
||
file_extension = file_extension.lower()
|
||
if file_extension in filelist: # filelist: wav.scp, file.jsonl;text.txt;
|
||
with open(data_in, encoding="utf-8") as fin:
|
||
for line in fin:
|
||
if not line.strip():
|
||
continue
|
||
key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
|
||
if data_in.endswith(".jsonl"): # file.jsonl: json.dumps({"source": data})
|
||
lines = json.loads(line.strip())
|
||
data = lines["source"]
|
||
key = lines.get("key", key)
|
||
else: # filelist, wav.scp, text.txt: id \t data or data
|
||
lines = line.strip().split(maxsplit=1)
|
||
data = lines[1] if len(lines) > 1 else lines[0]
|
||
key = lines[0] if len(lines) > 1 else key
|
||
|
||
data_list.append(data)
|
||
key_list.append(key)
|
||
else:
|
||
if key is None:
|
||
# key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
|
||
key = misc.extract_filename_without_extension(data_in)
|
||
data_list = [data_in]
|
||
key_list = [key]
|
||
elif isinstance(data_in, (list, tuple)):
|
||
if data_type is not None and isinstance(data_type, (list, tuple)): # mutiple inputs
|
||
data_list_tmp = []
|
||
for data_in_i, data_type_i in zip(data_in, data_type):
|
||
key_list, data_list_i = prepare_data_iterator(
|
||
data_in=data_in_i, data_type=data_type_i, input_format=input_format
|
||
)
|
||
data_list_tmp.append(data_list_i)
|
||
data_list = []
|
||
for item in zip(*data_list_tmp):
|
||
data_list.append(item)
|
||
else:
|
||
# [audio sample point, fbank, text]
|
||
data_list = data_in
|
||
if input_format != "auto":
|
||
data_list = [
|
||
load_bytes(item, input_format=input_format)
|
||
if isinstance(item, bytes) else item
|
||
for item in data_in
|
||
]
|
||
key_list = []
|
||
for data_i in data_in:
|
||
if isinstance(data_i, str) and os.path.exists(data_i):
|
||
key = misc.extract_filename_without_extension(data_i)
|
||
else:
|
||
if key is None:
|
||
key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
|
||
key_list.append(key)
|
||
|
||
else: # raw text; audio sample point, fbank; bytes
|
||
if isinstance(data_in, bytes): # audio bytes
|
||
data_in = load_bytes(data_in, input_format=input_format)
|
||
if key is None:
|
||
key = "rand_key_" + "".join(random.choice(chars) for _ in range(13))
|
||
data_list = [data_in]
|
||
key_list = [key]
|
||
|
||
return key_list, data_list
|
||
|
||
|
||
class AutoModel:
|
||
|
||
def __init__(self, **kwargs):
|
||
"""Initialize AutoModel with ASR model and optional sub-models.
|
||
|
||
Args:
|
||
model (str): Model name (hub alias or full ID) or local path.
|
||
device (str): Device for inference. "cuda:0", "cpu", "mps", "npu:0".
|
||
Logs a warning and falls back to CPU if unavailable (default: "cuda").
|
||
Explicit CPU use or ngpu=0 does not emit a fallback warning.
|
||
vad_model (str, optional): VAD model for long audio segmentation.
|
||
Enables processing of any-length audio.
|
||
vad_kwargs (dict, optional): VAD config, e.g. {"device": "cpu"}.
|
||
Explicit device overrides the ASR device; otherwise inherits its resolved device.
|
||
punc_model (str, optional): Punctuation restoration model.
|
||
Not needed for Fun-ASR-Nano/SenseVoice/Qwen3-ASR (they output punctuation natively).
|
||
punc_kwargs (dict, optional): Punctuation config; device follows the same rules as vad_kwargs.
|
||
spk_kwargs (dict, optional): Speaker config; device follows the same rules as vad_kwargs.
|
||
spk_model (str, optional): Speaker model for diarization ("cam++" or full model ID).
|
||
Requires vad_model. For Qwen3-ASR, also requires forced_aligner.
|
||
spk_mode (str, optional): Speaker diarization mode. "punc_segment" (default) or "vad_segment".
|
||
hub (str): Model hub. "ms" (ModelScope, default) or "hf" (HuggingFace).
|
||
ncpu (int): CPU threads (default: 4).
|
||
disable_update (bool): Skip version check on startup.
|
||
disable_pbar (bool): Disable tqdm progress bars.
|
||
**kwargs: Additional model-specific parameters (passed to config.yaml overrides).
|
||
|
||
Examples:
|
||
>>> model = AutoModel(model="paraformer-zh", vad_model="fsmn-vad", punc_model="ct-punc")
|
||
>>> model = AutoModel(model="FunAudioLLM/Fun-ASR-Nano-2512", trust_remote_code=True,
|
||
... remote_code="./model.py", vad_model="fsmn-vad", spk_model="cam++", hub="hf")
|
||
"""
|
||
if "vda_model" in kwargs:
|
||
raise TypeError(
|
||
"`vda_model` is not a valid AutoModel argument; use `vad_model` to enable voice activity detection."
|
||
)
|
||
|
||
try:
|
||
from funasr.utils.version_checker import check_for_update
|
||
|
||
check_for_update(disable=kwargs.get("disable_update", False))
|
||
except:
|
||
pass
|
||
|
||
log_level = getattr(logging, kwargs.get("log_level", "INFO").upper())
|
||
logging.basicConfig(level=log_level)
|
||
|
||
# Defaults and model construction must not mutate caller-owned submodel configs.
|
||
for name in ("vad_kwargs", "punc_kwargs", "spk_kwargs"):
|
||
kwargs[name] = copy.deepcopy(kwargs.get(name) or {})
|
||
|
||
model, kwargs = self.build_model(**kwargs)
|
||
|
||
# if vad_model is not None, build vad model else None
|
||
vad_model = kwargs.get("vad_model", None)
|
||
vad_kwargs = {} if kwargs.get("vad_kwargs", {}) is None else kwargs.get("vad_kwargs", {})
|
||
if vad_model is not None:
|
||
logging.info("Building VAD model.")
|
||
vad_kwargs["model"] = vad_model
|
||
vad_kwargs["model_revision"] = kwargs.get("vad_model_revision", "master")
|
||
vad_kwargs.setdefault("device", kwargs["device"])
|
||
vad_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4))
|
||
if "hub" in kwargs:
|
||
vad_kwargs.setdefault("hub", kwargs["hub"])
|
||
vad_model, vad_kwargs = self.build_model(**vad_kwargs)
|
||
|
||
# if punc_model is not None, build punc model else None
|
||
punc_model = kwargs.get("punc_model", None)
|
||
punc_kwargs = {} if kwargs.get("punc_kwargs", {}) is None else kwargs.get("punc_kwargs", {})
|
||
if punc_model is not None:
|
||
logging.info("Building punc model.")
|
||
punc_kwargs["model"] = punc_model
|
||
punc_kwargs["model_revision"] = kwargs.get("punc_model_revision", "master")
|
||
punc_kwargs.setdefault("device", kwargs["device"])
|
||
punc_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4))
|
||
if "hub" in kwargs:
|
||
punc_kwargs.setdefault("hub", kwargs["hub"])
|
||
punc_model, punc_kwargs = self.build_model(**punc_kwargs)
|
||
|
||
# if spk_model is not None, build spk model else None
|
||
spk_model = kwargs.get("spk_model", None)
|
||
spk_kwargs = {} if kwargs.get("spk_kwargs", {}) is None else kwargs.get("spk_kwargs", {})
|
||
cb_kwargs = (
|
||
{} if spk_kwargs.get("cb_kwargs", {}) is None else spk_kwargs.get("cb_kwargs", {})
|
||
)
|
||
if spk_model is not None:
|
||
logging.info("Building SPK model.")
|
||
spk_kwargs["model"] = spk_model
|
||
spk_kwargs["model_revision"] = kwargs.get("spk_model_revision", "master")
|
||
spk_kwargs.setdefault("device", kwargs["device"])
|
||
spk_kwargs.setdefault("ncpu", kwargs.get("ncpu", 4))
|
||
if "hub" in kwargs:
|
||
spk_kwargs.setdefault("hub", kwargs["hub"])
|
||
spk_model, spk_kwargs = self.build_model(**spk_kwargs)
|
||
self.cb_model = ClusterBackend(**cb_kwargs).to(kwargs["device"])
|
||
spk_mode = kwargs.get("spk_mode", "punc_segment")
|
||
if spk_mode not in ["default", "vad_segment", "punc_segment"]:
|
||
logging.error("spk_mode should be one of default, vad_segment and punc_segment.")
|
||
self.spk_mode = spk_mode
|
||
|
||
self.kwargs = kwargs
|
||
self.model = model
|
||
self.vad_model = vad_model
|
||
self.vad_kwargs = vad_kwargs
|
||
self.punc_model = punc_model
|
||
self.punc_kwargs = punc_kwargs
|
||
self.spk_model = spk_model
|
||
self.spk_kwargs = spk_kwargs
|
||
self.model_path = kwargs.get("model_path")
|
||
self._store_base_configs()
|
||
|
||
@staticmethod
|
||
def build_model(**kwargs):
|
||
"""Download model from hub, build all components, and load pretrained weights.
|
||
|
||
This method handles the full model construction pipeline:
|
||
1. Download model files from ModelScope/HuggingFace (if not local)
|
||
2. Parse config.yaml to determine model class, tokenizer, frontend
|
||
3. Instantiate tokenizer, frontend, and model via the registry
|
||
4. Load pretrained weights from model.pt
|
||
|
||
Args:
|
||
**kwargs: Must include 'model' (str). All other config.yaml fields can be overridden.
|
||
|
||
Returns:
|
||
tuple: (model, kwargs) where model is the instantiated nn.Module and
|
||
kwargs contains the resolved configuration.
|
||
"""
|
||
assert "model" in kwargs
|
||
# Silero VAD is loaded by its optional Python package rather than a
|
||
# FunASR model repository. Supplying model_conf keeps it on the normal
|
||
# AutoModel construction path while bypassing hub config resolution.
|
||
if kwargs["model"] in {"silero-vad", "silero_vad"}:
|
||
kwargs.setdefault("model_conf", {})
|
||
kwargs["model"] = "SileroVad"
|
||
if kwargs["model"] in {
|
||
"MOSS-Transcribe-Diarize",
|
||
"OpenMOSS-Team/MOSS-Transcribe-Diarize",
|
||
}:
|
||
kwargs.setdefault("model_conf", {})
|
||
kwargs.setdefault("model_path", kwargs["model"])
|
||
if "model_conf" not in kwargs:
|
||
logging.info("download models from model hub: {}".format(kwargs.get("hub", "ms")))
|
||
kwargs = download_model(**kwargs)
|
||
|
||
set_all_random_seed(kwargs.get("seed", 0))
|
||
|
||
device = kwargs.get("device", "cuda")
|
||
if (
|
||
(device.startswith("cuda") and not torch.cuda.is_available())
|
||
or (device.startswith("xpu") and not torch.xpu.is_available())
|
||
or (device.startswith("mps") and not torch.backends.mps.is_available())
|
||
or (device.startswith("npu") and not is_npu_available())
|
||
or kwargs.get("ngpu", 1) == 0
|
||
):
|
||
if kwargs.get("ngpu", 1) != 0:
|
||
logging.warning(
|
||
"AutoModel: requested device '%s' is unavailable for model '%s'; "
|
||
"falling back to CPU. Check your PyTorch build and accelerator "
|
||
"availability, or set device='cpu' to use CPU explicitly.",
|
||
device,
|
||
kwargs["model"],
|
||
)
|
||
device = "cpu"
|
||
kwargs["batch_size"] = 1
|
||
kwargs["device"] = device
|
||
|
||
ncpu = _resolve_ncpu(kwargs, 4)
|
||
kwargs["ncpu"] = ncpu
|
||
if torch.get_num_threads() != ncpu:
|
||
torch.set_num_threads(ncpu)
|
||
_limit_blas_threads(ncpu)
|
||
|
||
# build tokenizer
|
||
tokenizer = kwargs.get("tokenizer", None)
|
||
kwargs["tokenizer"] = tokenizer
|
||
kwargs["vocab_size"] = -1
|
||
|
||
if tokenizer is not None:
|
||
tokenizers = (
|
||
tokenizer.split(",") if isinstance(tokenizer, str) else tokenizer
|
||
) # type of tokenizers is list!!!
|
||
tokenizers_conf = kwargs.get("tokenizer_conf", {})
|
||
tokenizers_build = []
|
||
vocab_sizes = []
|
||
token_lists = []
|
||
|
||
### === only for kws ===
|
||
token_list_files = kwargs.get("token_lists", [])
|
||
seg_dicts = kwargs.get("seg_dicts", [])
|
||
### === only for kws ===
|
||
|
||
if not isinstance(tokenizers_conf, (list, tuple, ListConfig)):
|
||
tokenizers_conf = [tokenizers_conf] * len(tokenizers)
|
||
|
||
for i, tokenizer in enumerate(tokenizers):
|
||
tokenizer_class = tables.tokenizer_classes.get(tokenizer)
|
||
tokenizer_conf = tokenizers_conf[i]
|
||
|
||
### === only for kws ===
|
||
if len(token_list_files) > 1:
|
||
tokenizer_conf["token_list"] = token_list_files[i]
|
||
if len(seg_dicts) > 1:
|
||
tokenizer_conf["seg_dict"] = seg_dicts[i]
|
||
### === only for kws ===
|
||
|
||
tokenizer = tokenizer_class(**tokenizer_conf)
|
||
tokenizers_build.append(tokenizer)
|
||
token_list = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
|
||
token_list = (
|
||
tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else token_list
|
||
)
|
||
vocab_size = -1
|
||
if token_list is not None:
|
||
vocab_size = len(token_list)
|
||
|
||
if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"):
|
||
vocab_size = tokenizer.get_vocab_size()
|
||
token_lists.append(token_list)
|
||
vocab_sizes.append(vocab_size)
|
||
|
||
if len(tokenizers_build) <= 1:
|
||
tokenizers_build = tokenizers_build[0]
|
||
token_lists = token_lists[0]
|
||
vocab_sizes = vocab_sizes[0]
|
||
|
||
kwargs["tokenizer"] = tokenizers_build
|
||
kwargs["vocab_size"] = vocab_sizes
|
||
kwargs["token_list"] = token_lists
|
||
|
||
# build frontend
|
||
frontend = kwargs.get("frontend", None)
|
||
kwargs["input_size"] = None
|
||
if frontend is not None:
|
||
frontend_class = tables.frontend_classes.get(frontend)
|
||
frontend = frontend_class(**kwargs.get("frontend_conf", {}))
|
||
kwargs["input_size"] = (
|
||
frontend.output_size() if hasattr(frontend, "output_size") else None
|
||
)
|
||
kwargs["frontend"] = frontend
|
||
# build model
|
||
model_class = tables.model_classes.get(kwargs["model"])
|
||
if model_class is None:
|
||
raise RuntimeError(
|
||
_format_unregistered_component_error(
|
||
"model", kwargs["model"], tables.model_classes
|
||
)
|
||
)
|
||
model_conf = {}
|
||
deep_update(model_conf, kwargs.get("model_conf", {}))
|
||
deep_update(model_conf, kwargs)
|
||
model = model_class(**model_conf)
|
||
|
||
# init_param
|
||
init_param = kwargs.get("init_param", None)
|
||
if init_param is not None:
|
||
if os.path.exists(init_param):
|
||
logging.info(f"Loading pretrained params from {init_param}")
|
||
load_pretrained_model(
|
||
model=model,
|
||
path=init_param,
|
||
ignore_init_mismatch=kwargs.get("ignore_init_mismatch", True),
|
||
oss_bucket=kwargs.get("oss_bucket", None),
|
||
scope_map=kwargs.get("scope_map", []),
|
||
excludes=kwargs.get("excludes", None),
|
||
)
|
||
else:
|
||
print(f"error, init_param does not exist!: {init_param}")
|
||
|
||
# fp16
|
||
if kwargs.get("fp16", False):
|
||
model.to(torch.float16)
|
||
elif kwargs.get("bf16", False):
|
||
model.to(torch.bfloat16)
|
||
model.to(device)
|
||
model.eval()
|
||
|
||
if not kwargs.get("disable_log", True):
|
||
tables.print()
|
||
|
||
return model, kwargs
|
||
|
||
def __call__(self, *args, **cfg):
|
||
"""Internal: call .
|
||
|
||
Args:
|
||
*args: Variable positional arguments.
|
||
**cfg: Configuration overrides.
|
||
"""
|
||
kwargs = self.kwargs
|
||
self._merge_runtime_config(kwargs, cfg)
|
||
res = self.model(*args, kwargs)
|
||
return res
|
||
|
||
def generate(self, input, input_len=None, progress_callback=None, **cfg):
|
||
"""Run speech recognition on input audio.
|
||
|
||
This is the primary user-facing method. It automatically routes to:
|
||
- inference() if no vad_model is configured (single utterance)
|
||
- inference_with_vad() if vad_model is configured (long audio with segmentation)
|
||
|
||
Args:
|
||
input: Audio input. Accepts:
|
||
- File path (str): "audio.wav", "audio.mp3"
|
||
- URL (str): "https://..."
|
||
- numpy array: raw audio samples (float32, 16kHz)
|
||
- list: batch of file paths or arrays
|
||
- bytes: raw audio bytes
|
||
input_len (tensor, optional): Length of each input sample.
|
||
progress_callback (callable, optional): fn(current, total) called during processing.
|
||
**cfg: Runtime parameters:
|
||
- device: Placement is selected at construction. Runtime device overrides
|
||
are ignored; create another AutoModel to use a different device.
|
||
- cache (dict): State cache for streaming mode. Pass {} for first call.
|
||
- input_format (str): "auto" (default) or "pcm_s16le" to bypass
|
||
container detection for mono little-endian PCM16 byte inputs.
|
||
Use fs for the source sample rate.
|
||
- hotword (str/list): Keywords to boost recognition accuracy.
|
||
- postprocess_hotwords (str/list/dict): Text-level hotword correction after
|
||
decoding. Unlike model-level ``hotword``, this runs on the final text.
|
||
- postprocess_hotword_file (str): Hotword file path. Each line is a target
|
||
word or an explicit mapping like ``错误词=>目标词``.
|
||
- postprocess_hotword_threshold (float): Fuzzy match threshold in [0, 1].
|
||
- return_postprocess_hotword_matches (bool): Include replacement details.
|
||
- language (str): Language hint ("auto", "zh", "en", "Chinese", etc.)
|
||
- batch_size_s (int): Dynamic batch total duration in seconds.
|
||
- is_final (bool): Last chunk flag for streaming mode.
|
||
- return_spk_res (bool): Return speaker diarization results.
|
||
- sentence_timestamp (bool): Return sentence-level timestamps.
|
||
- use_itn (bool): Apply inverse text normalization (SenseVoice).
|
||
|
||
Returns:
|
||
list[dict]: Results for each input sample. Common fields:
|
||
- "key" (str): Sample identifier
|
||
- "text" (str): Recognized text
|
||
- "timestamp" (list): [[start_ms, end_ms], ...] per character/word
|
||
- "sentence_info" (list): [{text, start, end, spk, timestamp}, ...] when spk enabled
|
||
"""
|
||
self._reset_runtime_configs()
|
||
if self.vad_model is None:
|
||
results = self.inference(
|
||
input, input_len=input_len, progress_callback=progress_callback, **cfg
|
||
)
|
||
if self.punc_model is not None:
|
||
for result in results:
|
||
punc_res = self.inference(
|
||
result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg
|
||
)
|
||
if cfg.get("return_raw_text", self.kwargs.get("return_raw_text", False)):
|
||
result["raw_text"] = copy.copy(result["text"])
|
||
result["text"] = punc_res[0]["text"]
|
||
return apply_postprocess_hotwords_to_results(results, cfg)
|
||
|
||
else:
|
||
results = self.inference_with_vad(
|
||
input, input_len=input_len, progress_callback=progress_callback, **cfg
|
||
)
|
||
return apply_postprocess_hotwords_to_results(results, cfg)
|
||
|
||
def inference(
|
||
self,
|
||
input,
|
||
input_len=None,
|
||
model=None,
|
||
kwargs=None,
|
||
key=None,
|
||
progress_callback=None,
|
||
**cfg,
|
||
):
|
||
"""Run model inference on input data (internal method).
|
||
|
||
Handles batching, timing, and progress reporting. Called by generate()
|
||
and inference_with_vad(). Typically not called directly by users.
|
||
|
||
Args:
|
||
input: Audio data, file path, or text (for punc model).
|
||
input_len (tensor, optional): Input lengths for batch.
|
||
model (nn.Module, optional): Override model (used for VAD/PUNC/SPK sub-models).
|
||
kwargs (dict, optional): Override kwargs (used for sub-model configs).
|
||
key (list, optional): Sample identifiers.
|
||
progress_callback (callable, optional): Progress reporting function.
|
||
**cfg: Additional config merged into kwargs.
|
||
|
||
Returns:
|
||
list[dict]: Model inference results.
|
||
"""
|
||
if kwargs is None:
|
||
self._reset_runtime_configs()
|
||
kwargs = self.kwargs if kwargs is None else kwargs
|
||
if "cache" in kwargs:
|
||
kwargs.pop("cache")
|
||
self._merge_runtime_config(kwargs, cfg)
|
||
model = self.model if model is None else model
|
||
|
||
batch_size = kwargs.get("batch_size", 1)
|
||
# if kwargs.get("device", "cpu") == "cpu":
|
||
# batch_size = 1
|
||
|
||
key_list, data_list = prepare_data_iterator(
|
||
input, input_len=input_len, data_type=kwargs.get("data_type", None), key=key,
|
||
input_format=kwargs.get("input_format", "auto"),
|
||
)
|
||
|
||
speed_stats = {}
|
||
asr_result_list = []
|
||
num_samples = len(data_list)
|
||
disable_pbar = self.kwargs.get("disable_pbar", False)
|
||
pbar = (
|
||
tqdm(colour="blue", total=num_samples, dynamic_ncols=True) if not disable_pbar else None
|
||
)
|
||
time_speech_total = 0.0
|
||
time_escape_total = 0.0
|
||
for beg_idx in range(0, num_samples, batch_size):
|
||
end_idx = min(num_samples, beg_idx + batch_size)
|
||
data_batch = data_list[beg_idx:end_idx]
|
||
key_batch = key_list[beg_idx:end_idx]
|
||
batch = {"data_in": data_batch, "key": key_batch}
|
||
|
||
if (end_idx - beg_idx) == 1 and kwargs.get("data_type", None) == "fbank": # fbank
|
||
batch["data_in"] = data_batch[0]
|
||
batch["data_lengths"] = input_len
|
||
|
||
time1 = time.perf_counter()
|
||
with torch.no_grad():
|
||
res = model.inference(**batch, **kwargs)
|
||
if isinstance(res, (list, tuple)):
|
||
results = res[0] if len(res) > 0 else [{"text": ""}]
|
||
meta_data = res[1] if len(res) > 1 else {}
|
||
time2 = time.perf_counter()
|
||
|
||
asr_result_list.extend(results)
|
||
|
||
# batch_data_time = time_per_frame_s * data_batch_i["speech_lengths"].sum().item()
|
||
batch_data_time = meta_data.get("batch_data_time", -1)
|
||
time_escape = time2 - time1
|
||
speed_stats["load_data"] = meta_data.get("load_data", 0.0)
|
||
speed_stats["extract_feat"] = meta_data.get("extract_feat", 0.0)
|
||
speed_stats["forward"] = f"{time_escape:0.3f}"
|
||
speed_stats["batch_size"] = f"{len(results)}"
|
||
speed_stats["rtf"] = f"{(time_escape) / batch_data_time:0.3f}"
|
||
description = f"{speed_stats}, "
|
||
if pbar:
|
||
pbar.update(end_idx - beg_idx)
|
||
pbar.set_description(description)
|
||
if progress_callback:
|
||
try:
|
||
progress_callback(end_idx, num_samples)
|
||
except Exception as e:
|
||
logging.error(f"progress_callback error: {e}")
|
||
time_speech_total += batch_data_time
|
||
time_escape_total += time_escape
|
||
|
||
if pbar:
|
||
# pbar.update(1)
|
||
pbar.set_description(f"rtf_avg: {time_escape_total/time_speech_total:0.3f}")
|
||
|
||
device = next(model.parameters()).device
|
||
if device.type == "cuda":
|
||
with torch.cuda.device(device):
|
||
torch.cuda.empty_cache()
|
||
return asr_result_list
|
||
|
||
def inference_with_vad(self, input, input_len=None, **cfg):
|
||
"""Run ASR with VAD segmentation, punctuation, and optional speaker diarization.
|
||
|
||
Pipeline:
|
||
1. VAD: Segment audio into speech regions
|
||
2. ASR: Recognize each segment (sorted by length for efficient batching)
|
||
3. Timestamp merge: Combine per-segment timestamps with VAD offsets
|
||
4. Punctuation: Add punctuation to combined text (if punc_model configured)
|
||
5. Speaker diarization: Cluster speaker embeddings and assign labels (if spk_model configured)
|
||
|
||
Args:
|
||
input: Audio file path, URL, or numpy array.
|
||
input_len: Not used (kept for interface consistency).
|
||
**cfg: Runtime parameters (same as generate()).
|
||
|
||
Returns:
|
||
list[dict]: Results with fields: key, text, timestamp, sentence_info, raw_text.
|
||
"""
|
||
self._reset_runtime_configs()
|
||
if self.spk_model is not None and "output_timestamp" not in cfg:
|
||
cfg["output_timestamp"] = True
|
||
cfg["return_time_stamps"] = True
|
||
kwargs = self.kwargs
|
||
# step.1: compute the vad model
|
||
beg_vad = time.time()
|
||
vad_cfg = {
|
||
"input_format": kwargs.get("input_format", "auto"),
|
||
"fs": kwargs.get("fs", 16000),
|
||
**cfg,
|
||
}
|
||
res = self.inference(
|
||
input, input_len=input_len, model=self.vad_model, kwargs=self.vad_kwargs, **vad_cfg
|
||
)
|
||
end_vad = time.time()
|
||
|
||
# FIX(gcf): concat the vad clips for sense vocie model for better aed
|
||
if cfg.get("merge_vad", False):
|
||
for i in range(len(res)):
|
||
res[i]["value"] = merge_vad(
|
||
res[i]["value"], kwargs.get("merge_length_s", 15) * 1000
|
||
)
|
||
|
||
# step.2 compute asr model
|
||
model = self.model
|
||
self._merge_runtime_config(kwargs, cfg)
|
||
batch_size = max(int(kwargs.get("batch_size_s", 300)) * 1000, 1)
|
||
batch_size_threshold_ms = int(kwargs.get("batch_size_threshold_s", 60)) * 1000
|
||
kwargs["batch_size"] = batch_size
|
||
|
||
key_list, data_list = prepare_data_iterator(
|
||
input, input_len=input_len, data_type=kwargs.get("data_type", None),
|
||
input_format=kwargs.get("input_format", "auto"),
|
||
)
|
||
results_ret_list = []
|
||
# ASR updates kwargs with the resampled rate; retain the original for each input.
|
||
input_fs = kwargs.get("fs", 16000)
|
||
time_speech_total_all_samples = 1e-6
|
||
|
||
beg_total = time.time()
|
||
pbar_total = (
|
||
tqdm(colour="red", total=len(res), dynamic_ncols=True)
|
||
if not kwargs.get("disable_pbar", False)
|
||
else None
|
||
)
|
||
for i in range(len(res)):
|
||
key = res[i]["key"]
|
||
vadsegments = res[i]["value"]
|
||
input_i = data_list[i]
|
||
fs = kwargs["frontend"].fs if hasattr(kwargs["frontend"], "fs") else 16000
|
||
speech = load_audio_text_image_video(input_i, fs=fs, audio_fs=input_fs)
|
||
speech_lengths = len(speech)
|
||
n = len(vadsegments)
|
||
data_with_index = [(vadsegments[i], i) for i in range(n)]
|
||
sorted_data = sorted(data_with_index, key=lambda x: x[0][1] - x[0][0])
|
||
results_sorted = []
|
||
|
||
if not len(sorted_data):
|
||
results_ret_list.append({"key": key, "text": "", "timestamp": []})
|
||
logging.info("decoding, utt: {}, empty speech".format(key))
|
||
continue
|
||
|
||
if len(sorted_data) > 0 and len(sorted_data[0]) > 0:
|
||
batch_size = max(batch_size, sorted_data[0][0][1] - sorted_data[0][0][0])
|
||
|
||
if kwargs["device"] == "cpu":
|
||
batch_size = 0
|
||
|
||
beg_idx = 0
|
||
beg_asr_total = time.time()
|
||
time_speech_total_per_sample = speech_lengths / 16000
|
||
time_speech_total_all_samples += time_speech_total_per_sample
|
||
|
||
# pbar_sample = tqdm(colour="blue", total=n, dynamic_ncols=True)
|
||
|
||
all_segments = []
|
||
max_len_in_batch = 0
|
||
end_idx = 1
|
||
for j, _ in enumerate(range(0, n)):
|
||
# pbar_sample.update(1)
|
||
sample_length = sorted_data[j][0][1] - sorted_data[j][0][0]
|
||
potential_batch_length = max(max_len_in_batch, sample_length) * (j + 1 - beg_idx)
|
||
# batch_size_ms_cum += sorted_data[j][0][1] - sorted_data[j][0][0]
|
||
if (
|
||
j < n - 1
|
||
and sample_length < batch_size_threshold_ms
|
||
and potential_batch_length < batch_size
|
||
):
|
||
max_len_in_batch = max(max_len_in_batch, sample_length)
|
||
end_idx += 1
|
||
continue
|
||
|
||
speech_j, speech_lengths_j = slice_padding_audio_samples(
|
||
speech, speech_lengths, sorted_data[beg_idx:end_idx]
|
||
)
|
||
results = self.inference(
|
||
speech_j, input_len=None, model=model, kwargs=kwargs, **{**cfg, "fs": fs}
|
||
)
|
||
if self.spk_model is not None:
|
||
# compose vad segments: [[start_time_sec, end_time_sec, speech], [...]]
|
||
for _b in range(len(speech_j)):
|
||
vad_segments = [
|
||
[
|
||
sorted_data[beg_idx:end_idx][_b][0][0] / 1000.0,
|
||
sorted_data[beg_idx:end_idx][_b][0][1] / 1000.0,
|
||
np.array(speech_j[_b]),
|
||
]
|
||
]
|
||
segments = sv_chunk(vad_segments)
|
||
all_segments.extend(segments)
|
||
speech_b = [i[2] for i in segments]
|
||
spk_res = self.inference(
|
||
speech_b,
|
||
input_len=None,
|
||
model=self.spk_model,
|
||
kwargs=self.spk_kwargs,
|
||
**{**cfg, "fs": fs},
|
||
)
|
||
spk_embs = torch.cat([r["spk_embedding"] for r in spk_res], dim=0)
|
||
results[_b]["spk_embedding"] = spk_embs
|
||
beg_idx = end_idx
|
||
end_idx += 1
|
||
max_len_in_batch = sample_length
|
||
if len(results) < 1:
|
||
continue
|
||
results_sorted.extend(results)
|
||
|
||
# end_asr_total = time.time()
|
||
# time_escape_total_per_sample = end_asr_total - beg_asr_total
|
||
# pbar_sample.update(1)
|
||
# pbar_sample.set_description(f"rtf_avg_per_sample: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, "
|
||
# f"time_speech_total_per_sample: {time_speech_total_per_sample: 0.3f}, "
|
||
# f"time_escape_total_per_sample: {time_escape_total_per_sample:0.3f}")
|
||
|
||
if len(results_sorted) != n:
|
||
results_ret_list.append({"key": key, "text": "", "timestamp": []})
|
||
logging.info("decoding, utt: {}, empty result".format(key))
|
||
continue
|
||
restored_data = [0] * n
|
||
for j in range(n):
|
||
index = sorted_data[j][1]
|
||
restored_data[index] = results_sorted[j]
|
||
result = {}
|
||
|
||
# results combine for texts, timestamps, speaker embeddings and others
|
||
# TODO: rewrite for clean code
|
||
for j in range(n):
|
||
for k, v in restored_data[j].items():
|
||
if k.startswith("timestamp"):
|
||
if k not in result:
|
||
result[k] = []
|
||
for t in restored_data[j][k]:
|
||
if isinstance(t, dict):
|
||
t["start_time"] = (
|
||
float(t["start_time"]) * 1000 + int(vadsegments[j][0])
|
||
) / 1000
|
||
t["end_time"] = (
|
||
float(t["end_time"]) * 1000 + int(vadsegments[j][0])
|
||
) / 1000
|
||
else:
|
||
t[0] = int(t[0]) + int(vadsegments[j][0])
|
||
t[1] = int(t[1]) + int(vadsegments[j][0])
|
||
result[k].extend(restored_data[j][k])
|
||
elif k == "spk_embedding":
|
||
if k not in result:
|
||
result[k] = restored_data[j][k]
|
||
else:
|
||
result[k] = torch.cat([result[k], restored_data[j][k]], dim=0)
|
||
elif "text" in k:
|
||
if k not in result:
|
||
result[k] = restored_data[j][k]
|
||
else:
|
||
result[k] += " " + restored_data[j][k]
|
||
else:
|
||
if k not in result:
|
||
result[k] = restored_data[j][k]
|
||
else:
|
||
result[k] += restored_data[j][k]
|
||
|
||
# Convert dict-format timestamps (Fun-ASR-Nano) to list-format for downstream compatibility
|
||
if "timestamps" in result and "timestamp" not in result:
|
||
result["timestamp"] = [
|
||
[int(t["start_time"] * 1000), int(t["end_time"] * 1000)]
|
||
for t in result["timestamps"]
|
||
]
|
||
|
||
if not len(result["text"].strip()):
|
||
continue
|
||
return_raw_text = kwargs.get("return_raw_text", False)
|
||
aligned_words = result.get("words")
|
||
aligned_timestamps = result.get("timestamp")
|
||
aligned_word_text = None
|
||
if (
|
||
isinstance(aligned_words, list)
|
||
and aligned_words
|
||
and all(isinstance(word, str) and word.strip() for word in aligned_words)
|
||
and isinstance(aligned_timestamps, list)
|
||
and len(aligned_words) == len(aligned_timestamps)
|
||
):
|
||
aligned_word_text = " ".join(aligned_words)
|
||
|
||
# step.3 compute punc model
|
||
raw_text = None
|
||
punc_input_text = None
|
||
punc_res = None
|
||
punc_array = None
|
||
if self.punc_model is not None and "timestamps" not in result:
|
||
raw_text = copy.copy(result["text"])
|
||
punc_input_text = _join_vad_texts(
|
||
item.get("text", "") for item in restored_data
|
||
)
|
||
punc_res = self.inference(
|
||
punc_input_text,
|
||
model=self.punc_model,
|
||
kwargs=self.punc_kwargs,
|
||
**cfg,
|
||
)
|
||
if return_raw_text:
|
||
result["raw_text"] = raw_text
|
||
punc_array = punc_res[0].get("punc_array")
|
||
punctuated_surface = None
|
||
if aligned_word_text is not None:
|
||
punctuated_surface = _punctuate_surface_text(
|
||
punc_input_text, punc_array, self.punc_model
|
||
)
|
||
result["text"] = punctuated_surface or punc_res[0]["text"]
|
||
|
||
timestamp_text = punc_input_text
|
||
sentence_timestamps = result.get("timestamp", [])
|
||
punc_alignment_failed = False
|
||
if punc_res is not None:
|
||
try:
|
||
punc_length = len(punc_array)
|
||
except TypeError:
|
||
punc_length = -1
|
||
punc_array = None
|
||
if aligned_word_text is not None and punc_length == len(aligned_words):
|
||
timestamp_text = aligned_word_text
|
||
elif aligned_word_text is not None and punc_length > 0:
|
||
merged_units = _merge_timestamp_units(
|
||
punc_input_text,
|
||
aligned_words,
|
||
aligned_timestamps,
|
||
punc_array,
|
||
self.punc_model,
|
||
)
|
||
if merged_units is not None:
|
||
timestamp_text, sentence_timestamps = merged_units
|
||
if punc_array is not None and punc_length != len(sentence_timestamps):
|
||
punc_array = None
|
||
punc_alignment_failed = True
|
||
surface_sentence_list = None
|
||
if aligned_word_text is not None and punc_array is not None:
|
||
surface_sentence_list = _timestamp_sentences_from_surface(
|
||
punc_input_text,
|
||
sentence_timestamps,
|
||
punc_array,
|
||
self.punc_model,
|
||
return_raw_text=return_raw_text,
|
||
)
|
||
|
||
# speaker embedding cluster after resorted
|
||
if self.spk_model is not None and kwargs.get("return_spk_res", True):
|
||
if raw_text is None and self.spk_mode == "punc_segment":
|
||
logging.warning("punc_model is missing, falling back to vad_segment mode for speaker diarization.")
|
||
self.spk_mode = "vad_segment"
|
||
elif raw_text is None:
|
||
logging.error("Missing punc_model, which is required by spk_model.")
|
||
all_segments = sorted(all_segments, key=lambda x: x[0])
|
||
spk_embedding = result["spk_embedding"]
|
||
labels = self.cb_model(
|
||
spk_embedding.cpu(), oracle_num=kwargs.get("preset_spk_num", None)
|
||
)
|
||
# del result['spk_embedding']
|
||
# postprocess expects np.ndarray embeddings (per its type hint).
|
||
spk_embedding_np = spk_embedding.detach().cpu().numpy()
|
||
if kwargs.get("return_spk_center", False):
|
||
sv_output, spk_center = postprocess(
|
||
all_segments, None, labels, spk_embedding_np, return_spk_center=True
|
||
)
|
||
# Per-speaker ERes2NetV2 centroids, indexed by the `spk` id in
|
||
# sentence_info. Kept on the result for downstream voiceprint use
|
||
# (the per-chunk spk_embedding below is still deleted to keep output small).
|
||
result["spk_embedding_center"] = spk_center
|
||
else:
|
||
sv_output = postprocess(all_segments, None, labels, spk_embedding_np)
|
||
if self.spk_mode == "punc_segment" and "timestamp" not in result and "timestamps" not in result:
|
||
logging.warning("No timestamps in ASR result (e.g. SenseVoice), falling back to vad_segment mode for speaker diarization.")
|
||
self.spk_mode = "vad_segment"
|
||
if self.spk_mode == "vad_segment": # recover sentence_list
|
||
sentence_list = []
|
||
for rest, vadsegment in zip(restored_data, vadsegments):
|
||
if "timestamp" in rest:
|
||
ts = rest["timestamp"]
|
||
elif "timestamps" in rest:
|
||
ts = [
|
||
[int(t["start_time"] * 1000), int(t["end_time"] * 1000)]
|
||
for t in rest["timestamps"]
|
||
]
|
||
else:
|
||
logging.error("No timestamp found in ASR result. Speaker diarization relies on timestamps.")
|
||
ts = []
|
||
sentence_list.append(
|
||
{
|
||
"start": vadsegment[0],
|
||
"end": vadsegment[1],
|
||
"text": rest["text"],
|
||
"sentence": rest["text"],
|
||
"timestamp": ts,
|
||
}
|
||
)
|
||
elif self.spk_mode == "punc_segment":
|
||
if "timestamp" not in result:
|
||
logging.error(
|
||
"Only 'iic/speech_paraformer-large-vad-punc_asr_nat-zh-cn-16k-common-vocab8404-pytorch' \
|
||
and 'iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch'\
|
||
can predict timestamp, and speaker diarization relies on timestamps."
|
||
)
|
||
if punc_res is None:
|
||
logging.error(
|
||
"Missing punc_model, which is required for punc_segment speaker diarization."
|
||
)
|
||
sentence_list = []
|
||
elif surface_sentence_list is not None:
|
||
sentence_list = surface_sentence_list
|
||
elif kwargs.get("en_post_proc", False):
|
||
sentence_list = timestamp_sentence_en(
|
||
punc_array,
|
||
sentence_timestamps,
|
||
timestamp_text,
|
||
return_raw_text=return_raw_text,
|
||
)
|
||
else:
|
||
sentence_list = timestamp_sentence(
|
||
punc_array,
|
||
sentence_timestamps,
|
||
timestamp_text,
|
||
return_raw_text=return_raw_text,
|
||
)
|
||
distribute_spk(sentence_list, sv_output)
|
||
result["sentence_info"] = sentence_list
|
||
elif kwargs.get("sentence_timestamp", False):
|
||
if not len(result["text"].strip()):
|
||
sentence_list = []
|
||
elif self.punc_model is None and punc_res is None and not sentence_timestamps:
|
||
sentence_list = _vad_segment_sentences(restored_data, vadsegments)
|
||
elif punc_res is None:
|
||
logging.warning(
|
||
"punc_model is required for sentence_timestamp, skipping sentence segmentation."
|
||
)
|
||
sentence_list = []
|
||
elif punc_alignment_failed:
|
||
logging.warning(
|
||
"punctuation timestamps could not be aligned, falling back to VAD segments."
|
||
)
|
||
sentence_list = _vad_segment_sentences(restored_data, vadsegments)
|
||
elif surface_sentence_list is not None:
|
||
sentence_list = surface_sentence_list
|
||
else:
|
||
if kwargs.get("en_post_proc", False):
|
||
sentence_list = timestamp_sentence_en(
|
||
punc_array,
|
||
sentence_timestamps,
|
||
timestamp_text,
|
||
return_raw_text=return_raw_text,
|
||
)
|
||
else:
|
||
sentence_list = timestamp_sentence(
|
||
punc_array,
|
||
sentence_timestamps,
|
||
timestamp_text,
|
||
return_raw_text=return_raw_text,
|
||
)
|
||
result["sentence_info"] = sentence_list
|
||
if "spk_embedding" in result:
|
||
del result["spk_embedding"]
|
||
|
||
result["key"] = key
|
||
results_ret_list.append(result)
|
||
end_asr_total = time.time()
|
||
time_escape_total_per_sample = end_asr_total - beg_asr_total
|
||
if pbar_total:
|
||
pbar_total.update(1)
|
||
pbar_total.set_description(
|
||
f"rtf_avg: {time_escape_total_per_sample / time_speech_total_per_sample:0.3f}, "
|
||
f"time_speech: {time_speech_total_per_sample: 0.3f}, "
|
||
f"time_escape: {time_escape_total_per_sample:0.3f}"
|
||
)
|
||
|
||
# end_total = time.time()
|
||
# time_escape_total_all_samples = end_total - beg_total
|
||
# print(f"rtf_avg_all: {time_escape_total_all_samples / time_speech_total_all_samples:0.3f}, "
|
||
# f"time_speech_all: {time_speech_total_all_samples: 0.3f}, "
|
||
# f"time_escape_all: {time_escape_total_all_samples:0.3f}")
|
||
return results_ret_list
|
||
|
||
def export(self, input=None, **cfg):
|
||
"""Export model to ONNX format.
|
||
|
||
Creates a deep copy of the model to isolate ONNX operator monkey-patching,
|
||
then runs torch.onnx.export. The original model remains usable after export.
|
||
|
||
Args:
|
||
input: Sample input for tracing (auto-generated if None).
|
||
**cfg: Export parameters:
|
||
- type (str): Export format, "onnx" (default).
|
||
- quantize (bool): Whether to quantize the model.
|
||
- device (str): Device for export.
|
||
|
||
Returns:
|
||
str: Path to the exported model directory.
|
||
"""
|
||
"""
|
||
|
||
:param input:
|
||
:param type:
|
||
:param quantize:
|
||
:param fallback_num:
|
||
:param calib_num:
|
||
:param opset_version:
|
||
:param cfg:
|
||
:return:
|
||
"""
|
||
|
||
device = cfg.get("device", "cpu")
|
||
|
||
# 对模型进行深拷贝，隔离 ONNX 算子替换（Monkey-patching）对原模型的破坏
|
||
# Implement deep copy of the model to isolate ONNX operator monkey-patching
|
||
# and prevent corruption of the original model
|
||
model = copy.deepcopy(self.model).to(device=device)
|
||
|
||
# 对配置参数进行深拷贝，隔离 deep_update 和 del 的引用污染
|
||
# Implement deep copy of configuration parameters to isolate reference pollution caused by deep_update and del.
|
||
kwargs = copy.deepcopy(self.kwargs)
|
||
|
||
deep_update(kwargs, cfg)
|
||
kwargs["device"] = device
|
||
|
||
# Safely delete keys that may cause issues during export
|
||
if "model" in kwargs:
|
||
del kwargs["model"]
|
||
|
||
model.eval()
|
||
|
||
type = kwargs.get("type", "onnx")
|
||
|
||
key_list, data_list = prepare_data_iterator(
|
||
input, input_len=None, data_type=kwargs.get("data_type", None), key=None,
|
||
input_format=kwargs.get("input_format", "auto"),
|
||
)
|
||
|
||
with torch.no_grad():
|
||
# 这里的导出操作只会魔改 model 副本，原实例的 self.model 依然是纯洁的 PyTorch 图
|
||
# This export operation only mutates the model copy;
|
||
# the original self.model instance remains an intact PyTorch graph.
|
||
export_dir = export_utils.export(model=model, data_in=data_list, **kwargs)
|
||
|
||
return export_dir
|
||
|
||
@staticmethod
|
||
def _merge_runtime_config(kwargs, cfg):
|
||
"""Merge inference options without changing an already-loaded model's placement."""
|
||
# A runtime device string does not move weights. Keep the resolved device,
|
||
# including unavailable-device fallback, for feature creation and transfer.
|
||
runtime_cfg = {key: value for key, value in cfg.items() if key != "device"}
|
||
deep_update(kwargs, runtime_cfg)
|
||
|
||
def _store_base_configs(self):
|
||
"""Snapshot base kwargs for all submodules to allow reset before inference."""
|
||
baseline = {}
|
||
for name in dir(self):
|
||
if not name.endswith("kwargs"):
|
||
continue
|
||
value = getattr(self, name, None)
|
||
if isinstance(value, dict):
|
||
baseline[name] = copy.deepcopy(value)
|
||
# include primary kwargs explicitly
|
||
baseline["kwargs"] = copy.deepcopy(self.kwargs)
|
||
self._base_kwargs_map = baseline
|
||
|
||
_IMMUTABLE_KWARGS_KEYS = frozenset([
|
||
"token_list", "tokenizer", "frontend", "model", "init_param", "model_path",
|
||
])
|
||
|
||
def _reset_runtime_configs(self):
|
||
"""Ensure runtime kwargs reset to baseline defaults before inference."""
|
||
base_map = getattr(self, "_base_kwargs_map", None)
|
||
if not base_map:
|
||
return
|
||
|
||
for name, base in base_map.items():
|
||
restored = {}
|
||
for k, v in base.items():
|
||
if k in self._IMMUTABLE_KWARGS_KEYS or not isinstance(v, (dict, list)):
|
||
restored[k] = v
|
||
else:
|
||
restored[k] = copy.deepcopy(v)
|
||
setattr(self, name, restored)
|
||
|
||
ncpu = _resolve_ncpu(self.kwargs, 4)
|
||
self.kwargs["ncpu"] = ncpu
|
||
for name, value in base_map.items():
|
||
if name == "kwargs":
|
||
continue
|
||
config = getattr(self, name, None)
|
||
if isinstance(config, dict):
|
||
config.setdefault("ncpu", ncpu)
|
||
if torch.get_num_threads() != ncpu:
|
||
torch.set_num_threads(ncpu)
|