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
579 lines
23 KiB
Python
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())
|