voice_audio_multi/tests/parakeet_tdt/parakeet_python_warm_bench.py
Fabio c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
Initial import: audio.cpp with Orange Pi documentation
2026-11-04 19:05:22 +01:00

579 lines
23 KiB
Python

from __future__ import annotations
import argparse
import importlib.util
import json
import time
import types
from contextlib import contextmanager
from pathlib import Path
import librosa
REPO_ROOT = Path(__file__).resolve().parents[2]
REFERENCE_DRIVER = REPO_ROOT / "reference" / "parakeet-tdt" / "tools" / "test_driver.py"
def load_reference_driver():
spec = importlib.util.spec_from_file_location("parakeet_test_driver", REFERENCE_DRIVER)
if spec is None or spec.loader is None:
raise RuntimeError(f"failed to load reference driver from {REFERENCE_DRIVER}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Warm repeated Python benchmark for the Parakeet TDT reference.")
parser.add_argument("--model", default="models/parakeet-tdt-0.6b-v3")
parser.add_argument("--audio", type=Path, default=Path("build/assets/parakeet/2086-149220-0033_5s.wav"))
parser.add_argument("--warmup-audio", type=Path, default=None)
parser.add_argument("--audio-sequence", default="")
parser.add_argument("--backend", choices=("cpu", "cuda"), default="cpu")
parser.add_argument("--run-mode", choices=("offline", "longform", "streaming"), default="offline")
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--threads", type=int, default=8)
parser.add_argument("--warmup", type=int, default=2)
parser.add_argument("--iterations", type=int, default=5)
parser.add_argument("--batch-size", type=int, default=1)
parser.add_argument("--chunk-secs", type=float, default=2.0)
parser.add_argument("--left-context-secs", type=float, default=10.0)
parser.add_argument("--right-context-secs", type=float, default=2.0)
parser.add_argument("--timestamps", action="store_true", default=True)
parser.add_argument("--no-timestamps", dest="timestamps", action="store_false")
return parser.parse_args()
def parse_audio_sequence(args: argparse.Namespace) -> list[Path]:
if not args.audio_sequence:
return [args.audio]
return [Path(item) for item in args.audio_sequence.split(",") if item]
def synchronize_if_needed(torch_module, backend: str) -> None:
if backend == "cuda":
torch_module.cuda.synchronize()
def extract_word_timestamps(timestamp_payload: object) -> list[dict[str, object]]:
if not isinstance(timestamp_payload, dict):
return []
words = timestamp_payload.get("word")
if not isinstance(words, list):
return []
output: list[dict[str, object]] = []
for item in words:
if not isinstance(item, dict):
continue
try:
start = float(item["start"])
end = float(item["end"])
word = str(item["word"])
except Exception:
continue
output.append(
{
"start_sample": int(round(start * 16000.0)),
"end_sample": int(round(end * 16000.0)),
"word": word,
"confidence": 0.0,
}
)
return output
def make_jsonable(value: object) -> object:
if isinstance(value, dict):
return {str(key): make_jsonable(val) for key, val in value.items()}
if isinstance(value, list):
return [make_jsonable(item) for item in value]
if isinstance(value, tuple):
return [make_jsonable(item) for item in value]
if hasattr(value, "item") and callable(value.item):
try:
return value.item()
except Exception:
pass
return value
class SegmentTimer:
def __init__(self, torch_module, backend: str):
self._torch = torch_module
self._backend = backend
self._active = False
self._metrics = {
"frontend_ms": 0.0,
"pre_encode_ms": 0.0,
"encoder_total_ms": 0.0,
"decoder_ms": 0.0,
"timestamps_ms": 0.0,
}
def reset(self) -> None:
for key in self._metrics:
self._metrics[key] = 0.0
def begin(self) -> None:
self.reset()
self._active = True
def end(self) -> dict[str, float]:
self._active = False
encoder_ms = max(0.0, self._metrics["encoder_total_ms"] - self._metrics["pre_encode_ms"])
return {
"parakeet.frontend_ms": self._metrics["frontend_ms"],
"parakeet.pre_encode_ms": self._metrics["pre_encode_ms"],
"parakeet.encoder_ms": encoder_ms,
"parakeet.decoder_ms": self._metrics["decoder_ms"],
"parakeet.timestamps_ms": self._metrics["timestamps_ms"],
}
@contextmanager
def measure(self, key: str):
if not self._active:
yield
return
synchronize_if_needed(self._torch, self._backend)
started = time.perf_counter()
try:
yield
finally:
synchronize_if_needed(self._torch, self._backend)
ended = time.perf_counter()
self._metrics[key] += (ended - started) * 1000.0
def install_segment_hooks(asr_model, torch_module, backend: str):
import nemo.collections.asr.models.rnnt_models as rnnt_models_module
timer = SegmentTimer(torch_module, backend)
preprocessor = asr_model.preprocessor
original_preprocessor_forward = preprocessor.forward
def preprocessor_forward_wrapped(self, *args, **kwargs):
with timer.measure("frontend_ms"):
return original_preprocessor_forward(*args, **kwargs)
preprocessor.forward = types.MethodType(preprocessor_forward_wrapped, preprocessor)
encoder = asr_model.encoder
original_encoder_forward = encoder.forward
def encoder_forward_wrapped(self, *args, **kwargs):
with timer.measure("encoder_total_ms"):
return original_encoder_forward(*args, **kwargs)
encoder.forward = types.MethodType(encoder_forward_wrapped, encoder)
pre_encode = getattr(encoder, "pre_encode", None)
if pre_encode is not None and hasattr(pre_encode, "forward"):
original_pre_encode_forward = pre_encode.forward
def pre_encode_forward_wrapped(self, *args, **kwargs):
with timer.measure("pre_encode_ms"):
return original_pre_encode_forward(*args, **kwargs)
pre_encode.forward = types.MethodType(pre_encode_forward_wrapped, pre_encode)
decoding = asr_model.decoding
decoding_class = decoding.__class__
original_decoder_predictions = decoding_class.rnnt_decoder_predictions_tensor
def decoder_predictions_wrapped(self, *args, **kwargs):
with timer.measure("decoder_ms"):
return original_decoder_predictions(self, *args, **kwargs)
decoding_class.rnnt_decoder_predictions_tensor = decoder_predictions_wrapped
original_process_timestamp_outputs = rnnt_models_module.process_timestamp_outputs
def process_timestamp_outputs_wrapped(*args, **kwargs):
with timer.measure("timestamps_ms"):
return original_process_timestamp_outputs(*args, **kwargs)
rnnt_models_module.process_timestamp_outputs = process_timestamp_outputs_wrapped
return timer
def make_divisible_by(num: int, factor: int) -> int:
return (num // factor) * factor
def prepare_streaming_model(asr_model, open_dict_module, args: argparse.Namespace):
from nemo.collections.asr.models import EncDecHybridRNNTCTCModel, EncDecRNNTModel
from nemo.collections.asr.parts.utils.transcribe_utils import get_inference_dtype
if not isinstance(asr_model, EncDecRNNTModel) and not isinstance(asr_model, EncDecHybridRNNTCTCModel):
raise ValueError("Streaming warm bench currently supports RNNT / Hybrid RNNT models only")
decoding_cfg = asr_model.cfg.decoding
with open_dict_module(decoding_cfg):
decoding_cfg.tdt_include_token_duration = args.timestamps
decoding_cfg.greedy.preserve_alignments = False
decoding_cfg.fused_batch_size = -1
decoding_cfg.beam.return_best_hypothesis = True
if isinstance(asr_model, EncDecRNNTModel):
asr_model.change_decoding_strategy(decoding_cfg)
elif hasattr(asr_model, "cur_decoder"):
asr_model.change_decoding_strategy(decoding_cfg, decoder_type="rnnt")
asr_model.preprocessor.featurizer.dither = 0.0
asr_model.preprocessor.featurizer.pad_to = 0
asr_model.eval()
asr_model.freeze()
asr_model.to(get_inference_dtype(None, device=asr_model.device))
model_cfg = asr_model.cfg
audio_sample_rate = int(model_cfg.preprocessor["sample_rate"])
feature_stride_sec = float(model_cfg.preprocessor["window_stride"])
features_per_sec = 1.0 / feature_stride_sec
encoder_subsampling_factor = int(asr_model.encoder.subsampling_factor)
features_frame2audio_samples = make_divisible_by(
int(audio_sample_rate * feature_stride_sec),
factor=encoder_subsampling_factor,
)
encoder_frame2audio_samples = features_frame2audio_samples * encoder_subsampling_factor
from nemo.collections.asr.parts.utils.streaming_utils import ContextSize
context_encoder_frames = ContextSize(
left=int(args.left_context_secs * features_per_sec / encoder_subsampling_factor),
chunk=int(args.chunk_secs * features_per_sec / encoder_subsampling_factor),
right=int(args.right_context_secs * features_per_sec / encoder_subsampling_factor),
)
context_samples = ContextSize(
left=context_encoder_frames.left * encoder_subsampling_factor * features_frame2audio_samples,
chunk=context_encoder_frames.chunk * encoder_subsampling_factor * features_frame2audio_samples,
right=context_encoder_frames.right * encoder_subsampling_factor * features_frame2audio_samples,
)
if asr_model.cfg.encoder.att_context_style == "chunked_limited_with_rc":
asr_model.encoder.set_default_att_context_size(
att_context_size=[
context_encoder_frames.left,
context_encoder_frames.chunk,
context_encoder_frames.right,
]
)
return {
"audio_sample_rate": audio_sample_rate,
"encoder_frame2audio_samples": encoder_frame2audio_samples,
"context_samples": context_samples,
}
def run_streaming_once_in_process(
asr_model,
args: argparse.Namespace,
torch_module,
prepared: dict[str, object],
audio_path: Path,
) -> dict[str, object]:
from nemo.collections.asr.parts.utils.rnnt_utils import batched_hyps_to_hypotheses
from nemo.collections.asr.parts.utils.streaming_utils import StreamingBatchedAudioBuffer
from nemo.collections.asr.parts.utils.timestamp_utils import process_timestamp_outputs
audio_sample_rate = int(prepared["audio_sample_rate"])
encoder_frame2audio_samples = int(prepared["encoder_frame2audio_samples"])
context_samples = prepared["context_samples"]
samples, _ = librosa.load(str(audio_path), sr=audio_sample_rate, mono=True)
audio_batch = torch_module.from_numpy(samples).unsqueeze(0).to(asr_model.device)
audio_batch_lengths = torch_module.tensor([audio_batch.shape[1]], dtype=torch_module.long, device=asr_model.device)
batch_size = int(audio_batch.shape[0])
decoding_computer = asr_model.decoding.decoding.decoding_computer
state = None
current_batched_hyps = None
left_sample = 0
right_sample = min(context_samples.chunk + context_samples.right, audio_batch.shape[1])
buffer = StreamingBatchedAudioBuffer(
batch_size=batch_size,
context_samples=context_samples,
dtype=audio_batch.dtype,
device=audio_batch.device,
)
rest_audio_lengths = audio_batch_lengths.clone()
while left_sample < audio_batch.shape[1]:
chunk_length = min(right_sample, audio_batch.shape[1]) - left_sample
is_last_chunk_batch = chunk_length >= rest_audio_lengths
is_last_chunk = right_sample >= audio_batch.shape[1]
chunk_lengths_batch = torch_module.where(
is_last_chunk_batch,
rest_audio_lengths,
torch_module.full_like(rest_audio_lengths, fill_value=chunk_length),
)
buffer.add_audio_batch_(
audio_batch[:, left_sample:right_sample],
audio_lengths=chunk_lengths_batch,
is_last_chunk=is_last_chunk,
is_last_chunk_batch=is_last_chunk_batch,
)
encoder_output, encoder_output_len = asr_model(
input_signal=buffer.samples,
input_signal_length=buffer.context_size_batch.total(),
)
encoder_output = encoder_output.transpose(1, 2)
encoder_context = buffer.context_size.subsample(factor=encoder_frame2audio_samples)
encoder_context_batch = buffer.context_size_batch.subsample(factor=encoder_frame2audio_samples)
encoder_output = encoder_output[:, encoder_context.left :]
chunk_batched_hyps, _, state = decoding_computer(
x=encoder_output,
out_len=torch_module.where(
is_last_chunk_batch,
encoder_output_len - encoder_context_batch.left,
encoder_context_batch.chunk,
),
prev_batched_state=state,
multi_biasing_ids=None,
)
if current_batched_hyps is None:
current_batched_hyps = chunk_batched_hyps
else:
current_batched_hyps.merge_(chunk_batched_hyps)
rest_audio_lengths -= chunk_lengths_batch
left_sample = right_sample
right_sample = min(right_sample + context_samples.chunk, audio_batch.shape[1])
if current_batched_hyps is None:
raise RuntimeError("streaming decode produced no hypotheses")
hyp = batched_hyps_to_hypotheses(current_batched_hyps, None, batch_size=batch_size)[0]
hyp.text = asr_model.tokenizer.ids_to_text(hyp.y_sequence.tolist())
timestamp_payload = None
if args.timestamps:
hyp = asr_model.decoding.compute_rnnt_timestamps(hyp)
processed = process_timestamp_outputs(
hyp,
subsampling_factor=asr_model.encoder.subsampling_factor,
window_stride=asr_model.cfg["preprocessor"]["window_stride"],
)
if isinstance(processed, list):
hyp = processed[0]
else:
hyp = processed
timestamp_payload = make_jsonable(getattr(hyp, "timestamp", None))
return {
"text": hyp.text,
"timestamp": timestamp_payload,
}
def main() -> int:
args = parse_args()
audio_sequence = parse_audio_sequence(args)
driver = load_reference_driver()
device_label = f"cuda:{args.device}" if args.backend == "cuda" else "cpu"
driver_args = argparse.Namespace(
model=args.model,
audio=args.audio,
backend=args.backend,
device=args.device,
mode="perf",
run_mode=args.run_mode,
trace_log=None,
timing_log=None,
batch_size=args.batch_size,
timestamps=args.timestamps,
timestamp_levels="char,word,segment",
self_attention_model="rel_pos_local_attn",
att_context_left=256,
att_context_right=256,
chunk_secs=2.0,
left_context_secs=10.0,
right_context_secs=2.0,
clean_groundtruth_text=False,
langid="en",
)
driver.configure_runtime(driver_args)
nemo_asr, open_dict, torch, _, _ = driver.import_runtime_modules()
torch.set_num_threads(args.threads)
if hasattr(torch, "set_num_interop_threads"):
try:
torch.set_num_interop_threads(1)
except RuntimeError:
pass
runs: list[float] = []
run_metrics: list[dict[str, float]] = []
text_output = ""
word_timestamps: list[dict[str, object]] = []
sequence_steps: list[dict[str, object]] = []
if args.run_mode == "streaming":
asr_model = driver.resolve_model(nemo_asr, driver_args.model)
device = driver.configure_device(torch, asr_model, driver_args.backend)
device_label = str(device)
prepared_streaming = prepare_streaming_model(asr_model, open_dict, args)
def run_streaming_once():
synchronize_if_needed(torch, args.backend)
started = time.perf_counter()
last_sequence_steps: list[dict[str, object]] = []
for audio_path in audio_sequence:
step_started = time.perf_counter()
result = run_streaming_once_in_process(asr_model, args, torch, prepared_streaming, audio_path)
step_ended = time.perf_counter()
step_metrics = {
"parakeet.transcribe_wall_ms": (step_ended - step_started) * 1000.0,
}
last_sequence_steps.append(
{
"audio": str(audio_path),
"text_output": str(result.get("text", "")),
"word_timestamps": extract_word_timestamps(result.get("timestamp")) if args.timestamps else [],
"diagnostics": {},
"metrics": step_metrics,
}
)
synchronize_if_needed(torch, args.backend)
ended = time.perf_counter()
metrics = {
"parakeet.transcribe_wall_ms": (ended - started) * 1000.0,
}
return last_sequence_steps, metrics
for _ in range(max(0, args.warmup)):
run_streaming_once()
for _ in range(max(1, args.iterations)):
last_sequence_steps, metrics = run_streaming_once()
run_metrics.append(metrics)
runs.append(metrics["parakeet.transcribe_wall_ms"])
sequence_steps = last_sequence_steps
if last_sequence_steps:
text_output = str(last_sequence_steps[-1]["text_output"])
word_timestamps = list(last_sequence_steps[-1]["word_timestamps"])
else:
asr_model = driver.resolve_model(nemo_asr, driver_args.model)
device = driver.configure_device(torch, asr_model, driver_args.backend)
device_label = str(device)
if driver_args.timestamps:
driver.configure_timestamps(asr_model, open_dict)
if driver_args.run_mode == "longform":
driver.configure_longform(asr_model, driver_args)
segment_timer = install_segment_hooks(asr_model, torch, args.backend)
warmup_audio = args.warmup_audio if args.warmup_audio is not None else audio_sequence[0]
for _ in range(max(0, args.warmup)):
synchronize_if_needed(torch, args.backend)
_ = asr_model.transcribe(
[str(warmup_audio)],
batch_size=args.batch_size,
return_hypotheses=args.timestamps,
timestamps=args.timestamps,
)
synchronize_if_needed(torch, args.backend)
for _ in range(max(1, args.iterations)):
run_total_metrics: dict[str, float] = {}
started = time.perf_counter()
last_sequence_steps = []
for audio_path in audio_sequence:
synchronize_if_needed(torch, args.backend)
segment_timer.begin()
step_started = time.perf_counter()
hypotheses = asr_model.transcribe(
[str(audio_path)],
batch_size=args.batch_size,
return_hypotheses=args.timestamps,
timestamps=args.timestamps,
)
synchronize_if_needed(torch, args.backend)
step_ended = time.perf_counter()
metrics = segment_timer.end()
metrics["parakeet.transcribe_wall_ms"] = (step_ended - step_started) * 1000.0
result = hypotheses[0]
step_text = driver.extract_text(result)
step_words = extract_word_timestamps(driver.extract_timestamps(result)) if args.timestamps else []
last_sequence_steps.append(
{
"audio": str(audio_path),
"text_output": step_text,
"word_timestamps": step_words,
"diagnostics": {},
"metrics": metrics,
}
)
for key, value in metrics.items():
run_total_metrics[key] = run_total_metrics.get(key, 0.0) + value
ended = time.perf_counter()
metrics = run_total_metrics
metrics["parakeet.transcribe_wall_ms"] = (ended - started) * 1000.0
run_metrics.append(metrics)
runs.append(metrics["parakeet.transcribe_wall_ms"])
sequence_steps = last_sequence_steps
text_output = last_sequence_steps[-1]["text_output"] if last_sequence_steps else ""
word_timestamps = last_sequence_steps[-1]["word_timestamps"] if last_sequence_steps else []
average_metrics: dict[str, float] = {}
if run_metrics:
for key in run_metrics[0]:
average_metrics[key] = sum(run[key] for run in run_metrics) / float(len(run_metrics))
print("family=parakeet_tdt_reference")
print(f"backend={args.backend}")
print(f"run_mode={args.run_mode}")
print(f"device={device_label}")
print(f"threads={args.threads}")
print(f"batch_size={args.batch_size}")
print(f"timestamps={'true' if args.timestamps else 'false'}")
print(f"warmup={args.warmup}")
print(f"iterations={args.iterations}")
print(f"text_output={text_output}")
print("word_timestamps=" + json.dumps(word_timestamps, separators=(",", ":")))
if len(sequence_steps) > 1:
print(f"sequence_steps={len(sequence_steps)}")
for idx, step in enumerate(sequence_steps):
print(f"sequence_step[{idx}].audio={step['audio']}")
print(f"sequence_step[{idx}].text_output={step['text_output']}")
for idx, metrics in enumerate(run_metrics, start=1):
print(f"run={idx}")
for key, value in metrics.items():
print(f"{key}={value:.6f}")
print("average")
for key, value in average_metrics.items():
print(f"{key}={value:.6f}")
print(
"summary_json="
+ json.dumps(
{
"family": "parakeet_tdt_reference",
"backend": args.backend,
"run_mode": args.run_mode,
"device": device_label,
"threads": args.threads,
"batch_size": args.batch_size,
"timestamps": args.timestamps,
"warmup": args.warmup,
"iterations": args.iterations,
"warmup_audio": str(warmup_audio) if args.run_mode != "streaming" else "",
"audio_sequence": [str(path) for path in audio_sequence],
"text_output": text_output,
"word_timestamps": word_timestamps,
"sequence_steps": sequence_steps,
"runs": run_metrics,
"average": average_metrics,
},
separators=(",", ":"),
)
)
return 0
if __name__ == "__main__":
raise SystemExit(main())