diff --git a/apps/aiperf_qual/config/suites/librispeech-test-clean-speech-20.yaml b/apps/aiperf_qual/config/suites/librispeech-test-clean-speech-20.yaml index 2c71ea0941..c833eed543 100644 --- a/apps/aiperf_qual/config/suites/librispeech-test-clean-speech-20.yaml +++ b/apps/aiperf_qual/config/suites/librispeech-test-clean-speech-20.yaml @@ -3,14 +3,16 @@ # LibriSpeech test-clean (AIPerf `librispeech`, pinned): 20 utterances in seeded order, one per speaker, as the user # turn of a speech-to-speech model; no gold: the spoken answer is compared with the native model's. +# Version 2 resamples the 16 kHz utterances once to the model's 24 kHz input for both sides. suite: librispeech-test-clean-speech-20 -version: 1 +version: 2 source: kind: aiperf_public dataset: librispeech label_column: text cluster_column: speaker_id one_per_cluster: true + audio_sample_rate: 24000 selection: {method: seeded, count: 20} fields: {audio: audio_path} request: diff --git a/apps/aiperf_qual/tests/test_output_postprocess.py b/apps/aiperf_qual/tests/test_output_postprocess.py new file mode 100644 index 0000000000..0566c0e1e7 --- /dev/null +++ b/apps/aiperf_qual/tests/test_output_postprocess.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import sys + +import pytest + +from trtmc_aiperf_qual.absolute import postprocess +from trtmc_aiperf_qual.config import Environment + + +def test_postprocess_adds_text_without_changing_timing_or_original_artifacts(tmp_path): + script = tmp_path / "transcribe.py" + script.write_text("import json,sys\nfrom pathlib import Path\n" + "d=json.loads(Path(sys.argv[2]).read_text())\n" + "for row in d.values(): row['text']='Option C'\n" + "Path(sys.argv[4]).write_text(json.dumps(d))\n") + observations = {0: {"artifact": "unchanged.npy", "sample_rate": 24000}} + item = {"postprocess": {"script": script.name}} + environment = Environment({"repo": str(tmp_path), "serve_python": "/unused-serving-interpreter"}) + postprocess(environment, item, observations, tmp_path / "run", python=sys.executable) + assert observations == {0: {"artifact": "unchanged.npy", "sample_rate": 24000, "text": "Option C"}} + assert json.loads((tmp_path / "run.postprocess-input.json").read_text()) == { + "0": {"artifact": "unchanged.npy", "sample_rate": 24000}} + + +@pytest.mark.parametrize("mutation", ["d.clear()", "d['0']['sample_rate']=16000"]) +def test_postprocess_rejects_lost_samples_or_modified_evidence(tmp_path, mutation): + script = tmp_path / "bad.py" + script.write_text("import json,sys\nfrom pathlib import Path\n" + "d=json.loads(Path(sys.argv[2]).read_text())\n" + mutation + "\n" + "Path(sys.argv[4]).write_text(json.dumps(d))\n") + environment = Environment({"repo": str(tmp_path), "serve_python": sys.executable}) + with pytest.raises(ValueError, match="postprocessing changed"): + postprocess(environment, {"postprocess": {"script": script.name}}, + {0: {"sample_rate": 24000}}, tmp_path / "run") diff --git a/apps/aiperf_qual/tests/test_spoken_choice.py b/apps/aiperf_qual/tests/test_spoken_choice.py new file mode 100644 index 0000000000..f53e57091d --- /dev/null +++ b/apps/aiperf_qual/tests/test_spoken_choice.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Spoken answer scoring must not reward silence or failed extraction.""" + +import pytest + +from trtmc_aiperf_qual.gold_metrics import spoken_choice + +GOLD = {"answer": "C", "prompt": "Predators eat\nA. lions\nB. humans\nC. bunnies\nD. grass\nSelect A, B, C, or D."} + + +@pytest.mark.parametrize("text", ["Option C, bunnies.", "C.", "The answer is C.", "They eat bunnies."]) +def test_recognizes_the_correct_spoken_answer(text): + assert spoken_choice(GOLD, {"text": text}) == (True, "C") + + +@pytest.mark.parametrize("text", ["", "Thank you for listening.", "A predator needs food.", + "Option A or option C.", "lions and bunnies"]) +def test_missing_or_ambiguous_answers_are_never_guessed(text): + assert spoken_choice(GOLD, {"text": text}) == (False, "") + + +def test_wrong_answers_remain_wrong(): + assert spoken_choice(GOLD, {"text": "Option D, grass."}) == (False, "D") + + +def test_an_explicit_answer_overrides_a_word_in_the_explanation(): + assert spoken_choice(GOLD, {"text": "Bunnies are animals. The answer is A. Okay."}) == (False, "A") + + +@pytest.mark.parametrize("text", ["I would say C, definitely.", "You would pick C. That's the answer.", + "You would pick C. It could also be a human.", "Yes. The answer is C. Okay.", + "C is the correct answer."]) +def test_answer_cues_work_inside_a_spoken_reply(text): + assert spoken_choice(GOLD, {"text": text}) == (True, "C") + + +@pytest.mark.parametrize("text", ["The answer is B. C is the correct answer.", + "They don't eat bunnies.", "Not option C."]) +def test_conflicting_answers_and_negated_mentions_are_not_correct(text): + assert spoken_choice(GOLD, {"text": text}) == (False, "") + + +def test_question_and_answer_labels_are_not_required_by_transcription(): + # The scorer, rather than the ASR stage, owns access to gold labels. + assert spoken_choice(GOLD, {"text": "bunnies", "audio_artifact": "output.wav"}) == (True, "C") diff --git a/apps/aiperf_qual/tests/test_suite_audio.py b/apps/aiperf_qual/tests/test_suite_audio.py new file mode 100644 index 0000000000..727e1eb034 --- /dev/null +++ b/apps/aiperf_qual/tests/test_suite_audio.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A suite resamples dataset audio once, so the native and TRTMC sides read the same samples.""" + +import base64 +import io + +import numpy as np +import pytest + +from trtmc_aiperf_qual.config import ConfigError +from trtmc_aiperf_qual.suites import _hf_rows, _wav_media + +soundfile = pytest.importorskip("soundfile") +pytest.importorskip("scipy") + + +def _encoded(audio, rate): + buffer = io.BytesIO() + soundfile.write(buffer, audio, rate, format="WAV", subtype="PCM_16") + return buffer.getvalue() + + +def _decoded(media): + assert media["$file"]["suffix"] == ".wav" + return soundfile.read(io.BytesIO(base64.b64decode(media["$file"]["b64"])), dtype="float32") + + +def test_audio_is_resampled_once_to_the_suite_rate(): + tone = 0.5 * np.sin(2 * np.pi * 440 * np.arange(16_000) / 16_000) + audio, rate = _decoded(_wav_media(_encoded(tone, 16_000), 24_000)) + assert rate == 24_000 + assert audio.shape == (24_000,) + expected = 0.5 * np.sin(2 * np.pi * 440 * np.arange(24_000) / 24_000) + np.testing.assert_allclose(audio[200:-200], expected[200:-200], atol=2e-3) + + +@pytest.mark.parametrize("sample_rate", [None, 16_000]) +def test_audio_keeps_its_samples_without_a_different_suite_rate(sample_rate): + source = (np.arange(-800, 800) * 16).astype(np.float32) / 32768 # exact PCM16 values + audio, rate = _decoded(_wav_media(_encoded(source, 16_000), sample_rate)) + assert rate == 16_000 + np.testing.assert_array_equal(audio, source) + + +@pytest.mark.parametrize("sample_rate", [0, -24_000, 24_000.0, "24000", True]) +def test_invalid_suite_rates_are_rejected_before_the_dataset_loads(sample_rate): + with pytest.raises(ConfigError, match="audio_sample_rate"): + _hf_rows("dataset", object(), "revision", {"audio_sample_rate": sample_rate}, {}) diff --git a/apps/aiperf_qual/trtmc_aiperf_qual/absolute.py b/apps/aiperf_qual/trtmc_aiperf_qual/absolute.py index c51861d7dc..b09de4f17f 100644 --- a/apps/aiperf_qual/trtmc_aiperf_qual/absolute.py +++ b/apps/aiperf_qual/trtmc_aiperf_qual/absolute.py @@ -295,6 +295,8 @@ def _suite_side(environment: Environment, service: Mapping[str, Any], model: Map body = json.loads(record["responses"][-1]["text"]) key = request_sha(record["payload"]["request"]) by_request[key] = body.get("trtmc_observation") or {} + if item.get("postprocess"): + by_request[key]["artifact_directory"] = str(Path(service["records"]).parent / "scratch" / body["id"]) timing_by_request[key] = {"model_call_ms": float((body.get("trtmc_timing") or {}).get("model_call_ms", 0)), "completion_tokens": by_request[key].get("output_tokens")} observations = {index: by_request[problem["request_sha"]] for index, problem in enumerate(problems) @@ -308,9 +310,36 @@ def _suite_side(environment: Environment, service: Mapping[str, Any], model: Map side["rejected"] = {"greedy": out_of_capacity} if failing: side["failed"] = {"greedy": failed_reason(failing, len(failing))} + if item.get("postprocess"): + from .services import reference_python + + postprocess(environment, item, observations, out, python=reference_python(environment, model)) return side +def postprocess(environment: Environment, item: Mapping[str, Any], observations: dict[int, Any], out: Path, *, + python: str | None = None) -> None: + """Augment saved outputs after profiling; a separate process releases its model on exit.""" + import subprocess + + config = item["postprocess"] + source, target = out.parent / f"{out.name}.postprocess-input.json", out.parent / f"{out.name}.postprocess-output.json" + source.write_text(json.dumps(observations)) + command = [str(environment.values.get("scoring_python") or python or environment["serve_python"]), + str(environment.path("repo") / config["script"]), "--input", str(source), "--output", str(target), + *config.get("arguments", [])] + with (out.parent / f"{out.name}.postprocess.log").open("w") as log: + subprocess.run(command, check=True, stdout=log, stderr=subprocess.STDOUT, + timeout=float(config.get("timeout_s", 14400))) + updated = {int(index): value for index, value in json.loads(target.read_text()).items()} + if set(updated) != set(observations): + raise ValueError("output postprocessing changed the sample identities") + for index, original in observations.items(): + if any(updated[index].get(key) != value for key, value in original.items()): + raise ValueError("output postprocessing changed original evidence") + observations.update(updated) + + def run_side(environment: Environment, service: Mapping[str, Any], model: Mapping[str, Any], item: Mapping[str, Any], problems: Sequence[Mapping[str, Any]], out: Path, *, capacity: bool = False) -> dict[str, Any]: @@ -782,7 +811,8 @@ def send() -> None: # the whole exchange, an error body included, inside the ca def keeps_artifacts(model: Mapping[str, Any]) -> bool: """A benchmark reads the output files the servers write (their scratch is then kept).""" - return any(item.get("metric") in gold_metrics.ARTIFACT_METRICS for item in model.get("absolute", [])) + return any(item.get("postprocess") or item.get("metric") in gold_metrics.ARTIFACT_METRICS + for item in model.get("absolute", [])) def run_native(environment: Environment, model: Mapping[str, Any], python: str, plans: Mapping[str, Sequence], diff --git a/apps/aiperf_qual/trtmc_aiperf_qual/gold_metrics.py b/apps/aiperf_qual/trtmc_aiperf_qual/gold_metrics.py index 64b1e5bebb..fa69846f02 100644 --- a/apps/aiperf_qual/trtmc_aiperf_qual/gold_metrics.py +++ b/apps/aiperf_qual/trtmc_aiperf_qual/gold_metrics.py @@ -41,6 +41,39 @@ def choice(gold: Any, observation: Mapping[str, Any] | None, task: str | None = return bool(answer) and answer == str(gold).strip().upper(), answer +def spoken_choice(gold: Any, observation: Mapping[str, Any] | None, + task: str | None = None) -> tuple[bool, str]: + """Score a transcribed MCQ reply by an explicit letter or unique complete option text. + + Missing, ambiguous and unparsable answers are wrong. Never guess a letter. + """ + text = _text(observation).strip().lower() + def affirmative(start: int) -> bool: + clause = re.split(r"[.,;:!?]", text[:start])[-1] + return not re.search(r"\b(?:not|never|don't|doesn't|isn't|aren't|can't|cannot|won't)\b", clause) + + patterns = ( + r"\b(?:option|choice)\s+([a-d])\b", + r"\b(?:the\s+)?(?:correct\s+)?answer\s+(?:is|would be)\s+([a-d])(?=\s*(?:[.,:;!?]|$))", + r"\b([a-d])\s+is\s+(?:the\s+)?(?:correct\s+)?answer\b", + r"\b(?:i|you|we)(?:'d| would)?\s+(?:say|choose|pick)\s+([a-d])(?=\s*(?:[.,:;!?]|$))", + ) + letters = {match.group(1) for pattern in patterns for match in re.finditer(pattern, text) + if affirmative(match.start())} + direct = re.fullmatch(r"(?:the (?:correct )?answer is\s+)?([a-d])[.!]?", text) + if direct: + letters.add(direct.group(1)) + if not letters: + for letter, option in re.findall(r"(?m)^([A-D])\.\s*(.+)$", gold["prompt"]): + option = " ".join(re.findall(r"\w+", option.lower())) + raw = re.escape(option).replace(r"\ ", r"\W+") + positive = any(affirmative(match.start()) for match in re.finditer(r"(?= 3 and positive: + letters.add(letter.lower()) + answer = next(iter(letters)).upper() if len(letters) == 1 else "" + return bool(answer) and answer == str(gold["answer"]).strip().upper(), answer + + def _squash(text: str) -> str: return re.sub(r"\s+", " ", text.strip().lower()) @@ -206,7 +239,7 @@ def box_iou50(gold: Any, observation: Mapping[str, Any] | None, task: str | None return _iou(predicted, target) >= 0.5, match.group(0) -BINARY: dict[str, Callable[..., tuple[bool, str]]] = {"choice": choice, "contains": contains, "code_pass": code_pass, +BINARY: dict[str, Callable[..., tuple[bool, str]]] = {"choice": choice, "spoken_choice": spoken_choice, "contains": contains, "code_pass": code_pass, "top1": top1, "box_iou50": box_iou50} diff --git a/apps/aiperf_qual/trtmc_aiperf_qual/suites.py b/apps/aiperf_qual/trtmc_aiperf_qual/suites.py index afdf37890e..1f68a34eef 100644 --- a/apps/aiperf_qual/trtmc_aiperf_qual/suites.py +++ b/apps/aiperf_qual/trtmc_aiperf_qual/suites.py @@ -136,7 +136,8 @@ def _aiperf_public_records(source: Mapping[str, Any], selection: Mapping[str, An def _hf_dataset_records(source: Mapping[str, Any], selection: Mapping[str, Any]) -> list[dict[str, Any]]: """Selected rows of a Hugging Face dataset at a pinned revision (``hf_dataset``): ``dataset``, ``subset``, ``split``, ``revision``, and the ``prompt_column`` / ``image_column`` / - ``audio_column`` / ``label_column`` it reads; ``columns`` copies further fields.""" + ``audio_column`` / ``label_column`` it reads; ``columns`` copies further fields; + ``audio_sample_rate`` resamples the audio once for both sides.""" from types import SimpleNamespace meta = SimpleNamespace(hf_dataset_name=require(source, "dataset", "source"), hf_subset=source.get("subset"), @@ -254,12 +255,15 @@ def _interior_point(polygon: Any, size: Sequence[int]) -> tuple[float, float]: def _hf_rows(name: str, meta: Any, revision: str, source: Mapping[str, Any], selection: Mapping[str, Any]) -> list[dict[str, Any]]: - """Only selected rows are decoded. Audio becomes WAV (the TRTMC worker reads WAV); images keep - their original encoded bytes.""" + """Only selected rows are decoded. Audio becomes WAV (the TRTMC worker reads WAV), at the source's + ``audio_sample_rate`` when it sets one; images keep their original encoded bytes.""" + sample_rate = source.get("audio_sample_rate") + if sample_rate is not None and (isinstance(sample_rate, bool) or not isinstance(sample_rate, int) + or sample_rate <= 0): + raise ConfigError(f"audio_sample_rate must be a positive integer, not {sample_rate!r}") import io import datasets - import soundfile # Streaming datasets (for example LibriSpeech, whose config spans ~30 GB of train splits) are read # from the requested split only; media stay encoded until a row is selected. # ``data_files`` limits the download to the split's shards (for example ImageNet's validation parquet). @@ -309,12 +313,11 @@ def _hf_rows(name: str, meta: Any, revision: str, source: Mapping[str, Any], if "image" in roles: record["image"] = _encoded_media(row[roles["image"]], ".png") if "audio" in roles: - audio, rate = soundfile.read(io.BytesIO(row[roles["audio"]]["bytes"]), dtype="float32") - buffer = io.BytesIO() - soundfile.write(buffer, audio, rate, format="WAV", subtype="PCM_16") - record["audio"] = {"$file": {"suffix": ".wav", "b64": base64.b64encode(buffer.getvalue()).decode()}} + record["audio"] = _wav_media(row[roles["audio"]]["bytes"], sample_rate) if label_column: record["label"] = row[label_column] + if source.get("label_fields"): + record["label"] = {name: row[column] for name, column in source["label_fields"].items()} if source.get("label_image"): record["label"] = {"png_b64": base64.b64encode(row[label_column]["bytes"]).decode()} if source.get("polygon_label") and "image" in roles: # an object mask as a polygon, with a point inside it @@ -347,6 +350,24 @@ def _encoded_media(value: Mapping[str, Any], default_suffix: str) -> dict[str, A return {"$file": {"suffix": suffix, "b64": base64.b64encode(data).decode()}} +def _wav_media(data: bytes, sample_rate: int | None) -> dict[str, Any]: + """Encoded audio as a PCM16 WAV file, polyphase-resampled once to ``sample_rate`` when given: both the + native and the TRTMC side then read the same samples instead of resampling them differently.""" + import numpy as np + import soundfile + + audio, rate = soundfile.read(io.BytesIO(data), dtype="float32") + if sample_rate and rate != sample_rate: + from scipy.signal import resample_poly + + divisor = math.gcd(rate, sample_rate) + audio = np.clip(resample_poly(audio, sample_rate // divisor, rate // divisor, axis=0), -1.0, 1.0) + rate = sample_rate + buffer = io.BytesIO() + soundfile.write(buffer, audio, rate, format="WAV", subtype="PCM_16") + return {"$file": {"suffix": ".wav", "b64": base64.b64encode(buffer.getvalue()).decode()}} + + def _field(record: Mapping[str, Any], name: str) -> Any: """A record field; dotted names reach into nested objects and lists (``media.0.path``).""" value: Any = record diff --git a/apps/aiperf_qual/trtmc_aiperf_qual/transcribe_outputs.py b/apps/aiperf_qual/trtmc_aiperf_qual/transcribe_outputs.py new file mode 100644 index 0000000000..1b0b776a55 --- /dev/null +++ b/apps/aiperf_qual/trtmc_aiperf_qual/transcribe_outputs.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Transcribe saved generated audio after benchmark timing, without consulting labels.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--revision", required=True) + parser.add_argument("--last-seconds", type=float, default=0.0, + help="evaluate only the final reply window; zero keeps the complete audio") + args = parser.parse_args() + if args.last_seconds < 0: + parser.error("--last-seconds must not be negative") + + import librosa + import soundfile + import torch + from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor + + processor = AutoProcessor.from_pretrained(args.model, revision=args.revision) + model = AutoModelForSpeechSeq2Seq.from_pretrained( + args.model, revision=args.revision, dtype=torch.float16, + attn_implementation="sdpa").to("cuda").eval() + rows = json.loads(args.input.read_text()) + for index, observation in rows.items(): + if observation.get("audio_artifact"): + path = Path(observation["audio_artifact"]) + if not path.is_absolute(): + path = Path(observation["artifact_directory"]) / path + audio, rate = soundfile.read(path, dtype="float32") + else: + path = Path(observation["artifact"]) + audio, rate = np.load(path, allow_pickle=False), int(observation["sample_rate"]) + audio = np.asarray(audio).reshape(-1) + window_start = max(0, audio.size - round(args.last_seconds * rate)) if args.last_seconds else 0 + finite = np.isfinite(audio).all() + audio = audio[window_start:] + if not finite: + text = "" + elif not audio.size or not np.any(audio): + text = "" + else: + if rate != 16000: + audio = librosa.resample(audio, orig_sr=rate, target_sr=16000) + inputs = processor(audio, sampling_rate=16000, return_tensors="pt", return_attention_mask=True, + truncation=False, padding="longest") + with torch.inference_mode(): + generated = model.generate( + inputs.input_features.to("cuda", dtype=torch.float16), + attention_mask=inputs.attention_mask.to("cuda"), language="english", task="transcribe", + do_sample=False, temperature=0.0, num_beams=1, max_new_tokens=256, + return_timestamps=inputs.input_features.shape[-1] > 3000, + no_speech_threshold=0.6, logprob_threshold=-1.0, compression_ratio_threshold=2.4) + text = processor.batch_decode(generated, skip_special_tokens=True)[0].strip() + observation.update(text=text, asr={"model": args.model, "revision": args.revision, + "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "window_start_sample": window_start, "source_sample_rate": rate, + "last_seconds": args.last_seconds}) + print(json.dumps({"index": index, "text": text}), flush=True) + args.output.write_text(json.dumps(rows, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/families/personaplex/model.py b/families/personaplex/model.py index f1a0e4b377..5feac429e8 100644 --- a/families/personaplex/model.py +++ b/families/personaplex/model.py @@ -1415,11 +1415,22 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: parallel = ParallelConfig(tp_size=int(request.tensor_parallel_size)) parallel.validate() max_length = int(request.max_sequence_length or 512) + max_length = max(max_length, int(config.raw.get("personaplex_benchmark_capacity", 0))) config.raw["_model_dir"] = str(model_dir) config.raw["_fp32_layers"] = tuple(request.fp32_layers) model = _PersonaPlexModel() weights = model.load_weights(str(model_dir), config) + system_prompt = None + if prompt_path := config.raw.get("personaplex_system_prompt"): + from .system_prompt import load + + prompt_file = Path(prompt_path) + if not prompt_file.is_absolute(): + prompt_file = model_dir / prompt_file + system_prompt = load(prompt_file, _component_precision(config, _TEMPORAL_COMPONENT, request.precision), + int(weights["_hidden_size"]), max_length) + writer.set_header(family="personaplex", task=request.task, backend=request.backend) if parallel.enabled: for rank in range(parallel.tp_size): @@ -1492,4 +1503,8 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: "delays": [0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1], "text_prompt_ids": [], } + if system_prompt is not None: + runtime["system_prompt"], embeddings = system_prompt + runtime["listen_before_answer"] = bool(config.raw.get("personaplex_listen_before_answer", False)) + writer.add_bytes("system_prompt.embeddings", embeddings) writer.add_json("runtime.json", runtime) diff --git a/families/personaplex/reference/adapter.py b/families/personaplex/reference/adapter.py index 0d2836b6ca..979696a4d4 100644 --- a/families/personaplex/reference/adapter.py +++ b/families/personaplex/reference/adapter.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json import sys from pathlib import Path from typing import Any, Mapping @@ -42,28 +43,72 @@ def __init__(self, spec: Any, host: Any) -> None: for streaming in (self.mimi, self.other_mimi, self.generator): streaming.streaming_forever(1) warmup(self.mimi, self.other_mimi, self.generator, spec.device, self.frame_size) + self.system_prompt = None + adapter_options = getattr(spec, "options", {}) + self.listen_before_answer = bool(adapter_options.get("listen_before_answer", False)) + if prompt_path := adapter_options.get("system_prompt"): + prompt_path = Path(prompt_path) + if not prompt_path.is_absolute(): + prompt_path = Path(sys.prefix) / prompt_path + with np.load(prompt_path, allow_pickle=False) as prompt: + self.system_prompt = json.loads(str(prompt["metadata"])) + voice_embeddings = prompt["embeddings"][:self.system_prompt["voice_frames"]].copy() + voice = Path(adapter_options["voice_prompt"]) + if not voice.is_absolute(): + voice = Path(sys.prefix) / voice + if (self.system_prompt["precision"] != spec.precision + or self.system_prompt["voice"] != voice.name): + raise self.host.Error("the reference system prompt does not match its precision or voice") + self.generator.load_voice_prompt_embeddings(str(voice)) + # Official voice assets store BF16 embeddings; match the requested model storage dtype. + self.generator.voice_prompt_embeddings = self.generator.voice_prompt_embeddings.to(spec.dtype) + actual_voice = self.generator.voice_prompt_embeddings.float().cpu().numpy().reshape(-1, 4096) + if not np.array_equal(actual_voice, voice_embeddings): + raise self.host.Error("the reference voice embeddings differ from the prepared prefix") + self.generator.text_prompt_tokens = self.system_prompt["text_token_ids"] + self.generator.audio_silence_frame_cnt = self.system_prompt["silence_frames"] def invoke(self, request: Mapping[str, Any], artifact_base: Path) -> Any: + import soundfile import torch from moshi.models.lm import _iterate_audio, encode_from_sphn frames = int(request.get("max_new_tokens", 300)) - # Mono at Mimi's rate (librosa resampling): the environment's sphn stand-in reads 24 kHz files only. + # Suites deliver audio at Mimi's rate (``audio_sample_rate``), so the TRTMC side reads the same samples; + # the reference never resamples (the official sphn loader runs its resampler even on 24 kHz input). audio_path = str(self.host.required(request, "audio_path")) + if (rate := soundfile.info(audio_path).samplerate) != self.mimi.sample_rate: + raise self.host.Error(f"PersonaPlex reference input must be {self.mimi.sample_rate} Hz audio, not {rate} Hz") source_audio = self.host.load_audio(audio_path, int(self.mimi.sample_rate))[None, :] - for streaming in (self.mimi, self.other_mimi, self.generator): - streaming.reset_streaming() - + question_frames = (source_audio.shape[-1] + self.frame_size - 1) // self.frame_size + system_prompt = getattr(self, "system_prompt", None) + if system_prompt is not None: + tail = int(request.get("tail_frames", 0)) + if tail < 0 or tail > frames: + raise self.host.Error("tail_frames must fit the output frame budget") + source_audio = np.pad(source_audio, ((0, 0), (0, tail * self.frame_size))) def run() -> list[Any]: + for streaming in (self.mimi, self.other_mimi, self.generator): + streaming.reset_streaming() + if system_prompt is not None: + self.generator.step_system_prompts(self.mimi) + self.mimi.reset_streaming() chunks: list[Any] = [] + input_frame = 0 for encoded in encode_from_sphn(self.mimi, _iterate_audio(source_audio, sample_interval_size=self.frame_size, pad=True), max_batch=1): for index in range(encoded.shape[-1]): - tokens = self.generator.step(encoded[:, :, index:index + 1]) + forced = ({"moshi_tokens": self.generator._encode_zero_frame(), + "text_token": self.generator.zero_text_code} + if getattr(self, "listen_before_answer", False) and input_frame < question_frames else {}) + if getattr(self, "listen_before_answer", False) and input_frame == question_frames: + forced = {"text_token": 0} # EPAD: official Moshi spoken-QA response trigger. + tokens = self.generator.step(encoded[:, :, index:index + 1], **forced) + input_frame += 1 if tokens is None: continue + # Decode the returned stream once; the other codec is only needed for official warmup. chunks.append(self.mimi.decode(tokens[:, 1:9])[0, 0].float()) - self.other_mimi.decode(tokens[:, 1:9]) if len(chunks) >= frames: return chunks return chunks @@ -73,5 +118,9 @@ def run() -> list[Any]: raise self.host.Error("the native PersonaPlex produced no speech frames") audio = torch.cat(chunks).cpu().numpy().astype(np.float32) rate = int(self.mimi.sample_rate) + provenance = ({"system_prompt_voice": system_prompt["voice"], + "system_prompt_frames": system_prompt["frames"]} + if system_prompt is not None else {}) return self.host.invocation({**self.host.tensor_observation(audio, artifact_base), "sample_rate": rate, - "output_tokens": len(chunks), "audio_seconds": audio.size / rate}, model_ms) + "output_tokens": len(chunks), "audio_seconds": audio.size / rate, + **provenance}, model_ms) diff --git a/families/personaplex/reference/prepare_voicebench.py b/families/personaplex/reference/prepare_voicebench.py new file mode 100644 index 0000000000..96aa507735 --- /dev/null +++ b/families/personaplex/reference/prepare_voicebench.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prepare a small source overlay and official prompt for the VoiceBench profile. + +Checkpoints stay in the shared HF cache. The environment contains symlinks, +one modified config, the selected voice and a derived prompt artifact only. +""" + +from __future__ import annotations + +import io +import json +import math +from pathlib import Path +import sys +import tarfile + +from datasets import Audio, load_dataset +from huggingface_hub import snapshot_download +import soundfile + +from families.personaplex.reference.prepare import main as prepare_official +from families.personaplex.system_prompt import prepare + +MODEL_REVISION = "fdaf4090a61cb315c138a1faee287ffd6c716309" +DATASET_REVISION = "b02edcef1330480be3a11bd6f7434ac32f05ad08" +TEXT_PROMPT = ("You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way. " + "For multiple-choice questions, say the option letter and its answer. Keep your answer brief.") + + +def main() -> None: + prepare_official() + source = Path(snapshot_download("nvidia/personaplex-7b-v1", revision=MODEL_REVISION)) + target = Path(sys.prefix) / "trtmc-reference/personaplex-model" + target.mkdir(parents=True, exist_ok=True) + for p in source.iterdir(): + if p.name != "config.json" and not (target / p.name).exists(): + (target / p.name).symlink_to(p) + voice = target / "NATM1.pt" + with tarfile.open(source / "voices.tgz") as archive: + voice.write_bytes(archive.extractfile("voices/NATM1.pt").read()) + prompt = target / "system_prompt.npz" + metadata = prepare(source, voice, TEXT_PROMPT, "fp16", prompt) + data = load_dataset("hlt-lab/voicebench", "openbookqa", split="test", revision=DATASET_REVISION) + data = data.cast_column("audio", Audio(decode=False)) + if len(data) != 455: + raise ValueError("the pinned VoiceBench OpenBookQA split must contain 455 questions") + durations = [soundfile.info(io.BytesIO(row["audio"]["bytes"])).duration for row in data] + required_frames = metadata["frames"] + math.ceil(max(durations) * 12.5) + 125 + capacity = math.ceil(required_frames / 128) * 128 + config = json.loads((source / "config.json").read_text()) + config.update(personaplex_system_prompt="system_prompt.npz", personaplex_benchmark_capacity=capacity, + personaplex_listen_before_answer=True) + (target / "config.json").write_text(json.dumps(config, indent=2) + "\n") + (target / "benchmark-capacity.json").write_text(json.dumps({"dataset_revision": DATASET_REVISION, + "samples": len(data), "max_question_seconds": max(durations), "prompt_frames": metadata["frames"], + "reply_frames": 125, "required_frames": required_frames, "capacity": capacity}, indent=2) + "\n") + print(json.dumps({"source": str(target), "prompt": metadata, "capacity": capacity})) + + +if __name__ == "__main__": + main() diff --git a/families/personaplex/runtime/pipeline.cpp b/families/personaplex/runtime/pipeline.cpp index afd548414f..1408b4277d 100644 --- a/families/personaplex/runtime/pipeline.cpp +++ b/families/personaplex/runtime/pipeline.cpp @@ -756,11 +756,20 @@ void SpeechPipeline::speak_run_generation_loop(const SpeechGenerationSettings& s CudaStageTimer depth_timer; frames_collected = 0; - for (int32_t offset = 0; offset < plan.total_iters && frames_collected < plan.output_frames; - ++offset) { + for (int32_t step = 0; step < plan.total_iters && frames_collected < plan.output_frames; + ++step) { + const int32_t offset = config_.system_prompt_offset + step; write_user_tokens_to_delay_cache(delay_state, codec_tokens, offset, settings.stream_cb, settings.num_frames, settings.encode_codebooks, - settings.audio_bos); + settings.audio_bos, step); + if (step < settings.listen_frames) + write_listening_tokens_to_delay_cache(delay_state, offset, settings.text_pad_id); + else if (config_.listen_before_answer && step == settings.listen_frames) { + // EPAD (0) starts the response after the complete spoken question. + const auto index = delay_cache_index(delay_state, 0, offset); + delay_state.cache[index] = 0; + delay_state.provided[index] = 1; + } fill_initial_delay_tokens(delay_state, offset, settings.text_bos, settings.audio_bos); if (offset == 0) { seed_delay_offset_zero(delay_state, settings.text_bos, settings.audio_bos); @@ -804,6 +813,12 @@ void SpeechPipeline::speak_run_generation_loop(const SpeechGenerationSettings& s ++frames_collected; } + if (config_.system_prompt_offset > 0) { + // The benchmark's input includes its reply window; process the same timeline on both + // sides. + speech_maybe_log_interleaved_debug(offset, text_input, sampled_text_token, frame_codes); + continue; + } SpeechDecodeStopInput stop_input; stop_input.text_eos_token_id = config_.text_eos_token_id; stop_input.text_padding_id = config_.text_padding_id; @@ -864,6 +879,22 @@ AudioResult SpeechPipeline::speak(const float* audio_in, int32_t num_samples, input_sample_rate = target_rate; } + std::vector reply_window; + if (config_.system_prompt_offset > 0 && cfg.tail_frames != 0) { + if (cfg.tail_frames < 0 || cfg.tail_frames > max_output_frames) + throw std::invalid_argument("tail_frames must fit the output frame budget"); + const auto frame_size = + static_cast(std::lround(config_.sample_rate / config_.frame_rate)); + const int64_t padded = static_cast(samples_count) + + static_cast(cfg.tail_frames) * frame_size; + if (padded > std::numeric_limits::max()) + throw std::invalid_argument("PersonaPlex reply window is too large"); + reply_window.assign(samples_ptr, samples_ptr + samples_count); + reply_window.resize(static_cast(padded), 0.0F); + samples_ptr = reply_window.data(); + samples_count = static_cast(padded); + } + std::cerr << "[speech] Starting pipeline with " << samples_count << " input samples" << std::endl; @@ -892,12 +923,26 @@ AudioResult SpeechPipeline::speak(const float* audio_in, int32_t num_samples, temporal_state_->reset(); - if (should_run_text_prompt_injection(config_)) + if (!config_.system_prompt_embeddings.empty()) { + const int32_t hidden = config_.temporal_hidden_size; + for (std::size_t index = 0; index < config_.system_prompt_embeddings.size(); + index += hidden) + run_temporal_embed_step(config_.system_prompt_embeddings.data() + index, hidden); + } else if (should_run_text_prompt_injection(config_)) run_text_prompt(); const int32_t num_cb = config_.num_codebooks; const int32_t hidden = config_.temporal_hidden_size; auto delay_state = make_delay_cache_state(config_.delays, num_cb); + if (config_.system_prompt_offset > 0) { + if (config_.system_prompt_cache.size() != delay_state.cache.size() || + config_.system_prompt_provided.size() != delay_state.provided.size() || + config_.system_prompt_offset - 1 + num_frames > temporal_state_->max_length()) + throw std::invalid_argument( + "PersonaPlex prompt and audio exceed the bundle context or cache geometry"); + delay_state.cache = config_.system_prompt_cache; + delay_state.provided = config_.system_prompt_provided; + } SpeechOutputPlanInput plan_input; plan_input.sample_rate = config_.sample_rate; plan_input.frame_rate = config_.frame_rate; @@ -907,7 +952,13 @@ AudioResult SpeechPipeline::speak(const float* audio_in, int32_t num_samples, plan_input.tail_frames = cfg.tail_frames; plan_input.max_output_frames = max_output_frames; plan_input.max_delay = delay_state.max_delay; - const auto plan = ComputeSpeechOutputPlan(plan_input); + auto plan = ComputeSpeechOutputPlan(plan_input); + if (config_.system_prompt_offset > 0) { + // A primed LMGen returns one frame per user step; it has no cold-start delay. + plan.effective_frames = num_frames; + plan.output_frames = std::min(num_frames, max_output_frames); + plan.total_iters = plan.output_frames; + } const int32_t mimi_cb = config_.mimi_decode_codebooks; std::vector output_codes; output_codes.reserve(static_cast(mimi_cb) * plan.output_frames); @@ -918,8 +969,10 @@ AudioResult SpeechPipeline::speak(const float* audio_in, int32_t num_samples, << ", input_effective=" << plan.effective_frames << ", tail_frames=" << plan.extra_tail << ")" << std::endl; - const SpeechGenerationSettings settings = + SpeechGenerationSettings settings = make_speech_generation_settings(config_, hidden, encoder_shape); + if (config_.listen_before_answer) + settings.listen_frames = std::max(0, num_frames - cfg.tail_frames); int32_t frames_collected = 0; speak_run_generation_loop(settings, plan, delay_state, codec_tokens, output_codes, diff --git a/families/personaplex/runtime/plugin.cpp b/families/personaplex/runtime/plugin.cpp index bc11a8db57..772f514304 100644 --- a/families/personaplex/runtime/plugin.cpp +++ b/families/personaplex/runtime/plugin.cpp @@ -55,6 +55,22 @@ SpeechConfig parse_config(const BundleReader& bundle, const nlohmann::json& json config.temporal_text_embedding = floats("temporal_text.embeddings"); config.depth_text_embedding = floats("depth_text.embeddings"); config.depth_audio_embeddings = floats("depth_audio.embeddings"); + if (json.contains("system_prompt")) { + config.listen_before_answer = json.value("listen_before_answer", false); + const auto& prompt = json.at("system_prompt"); + config.system_prompt_offset = prompt.at("offset").get(); + config.system_prompt_cache = prompt.at("cache").get>(); + config.system_prompt_provided = prompt.at("provided").get>(); + config.system_prompt_embeddings = floats("system_prompt.embeddings"); + const auto frames = prompt.at("frames").get(); + if (prompt.at("version").get() != 1 || frames <= 0 || + config.system_prompt_offset != frames + 1 || + config.system_prompt_embeddings.size() != + static_cast(frames) * config.temporal_hidden_size || + config.system_prompt_cache.size() != 68 || config.system_prompt_provided.size() != 68 || + !config.text_prompt_ids.empty()) + throw std::runtime_error("PersonaPlex system-prompt state is invalid"); + } if (config.sample_rate <= 0 || config.num_codebooks <= 0 || config.codebook_size <= 0 || config.temporal_hidden_size <= 0 || config.temporal_num_layers <= 0 || config.depth_hidden_size <= 0 || config.depth_num_layers <= 0 || diff --git a/families/personaplex/runtime/speech_config.h b/families/personaplex/runtime/speech_config.h index bbe7247193..9b9069c9d5 100644 --- a/families/personaplex/runtime/speech_config.h +++ b/families/personaplex/runtime/speech_config.h @@ -54,6 +54,11 @@ struct SpeechConfig { int32_t text_eos_token_id{-1}; std::vector text_prompt_ids; + std::vector system_prompt_embeddings; + std::vector system_prompt_cache; + std::vector system_prompt_provided; + int32_t system_prompt_offset{0}; + bool listen_before_answer{false}; }; } // namespace trtmc diff --git a/families/personaplex/runtime/speech_delay_cache.h b/families/personaplex/runtime/speech_delay_cache.h index 26691777b1..81761bd529 100644 --- a/families/personaplex/runtime/speech_delay_cache.h +++ b/families/personaplex/runtime/speech_delay_cache.h @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include #include @@ -67,12 +68,14 @@ inline DelayCacheState make_delay_cache_state(const std::vector& config inline void write_user_tokens_to_delay_cache(DelayCacheState& delay_state, const std::vector& codec_tokens, int32_t offset, int32_t stream_cb, int32_t num_frames, - int32_t encode_codebooks, int32_t audio_bos) { + int32_t encode_codebooks, int32_t audio_bos, + int32_t input_frame = -1) { + const int32_t frame = input_frame >= 0 ? input_frame : offset; for (int32_t cb = 0; cb < stream_cb; ++cb) { const int32_t k = stream_cb + 1 + cb; int32_t user_tok = audio_bos; - if (offset < num_frames) { - const auto tok_idx = static_cast(offset) * encode_codebooks + cb; + if (frame < num_frames) { + const auto tok_idx = static_cast(frame) * encode_codebooks + cb; if (tok_idx < codec_tokens.size()) { user_tok = codec_tokens[tok_idx]; } @@ -97,6 +100,17 @@ inline void fill_initial_delay_tokens(DelayCacheState& delay_state, int32_t offs } } +inline void write_listening_tokens_to_delay_cache(DelayCacheState& state, int32_t offset, + int32_t text_padding_id) { + // The same forced agent-silence codes used by official LMGen system prompts. + constexpr std::array silence{948, 243, 1178, 546, 1736, 1030, 1978, 2008}; + for (int32_t k = 0; k <= 8; ++k) { + const auto index = delay_cache_index(state, k, offset + state.delays[k]); + state.cache[index] = k == 0 ? text_padding_id : silence[k - 1]; + state.provided[index] = 1; + } +} + inline void seed_delay_offset_zero(DelayCacheState& delay_state, int32_t text_bos, int32_t audio_bos) { for (int32_t k = 0; k < delay_state.total_k; ++k) { diff --git a/families/personaplex/runtime/speech_runtime_plan.h b/families/personaplex/runtime/speech_runtime_plan.h index 27c493053c..80245cc8ab 100644 --- a/families/personaplex/runtime/speech_runtime_plan.h +++ b/families/personaplex/runtime/speech_runtime_plan.h @@ -20,6 +20,7 @@ struct EncoderShapeInfo { }; struct SpeechGenerationSettings { + int32_t listen_frames{0}; int32_t hidden{0}; int32_t num_cb{0}; int32_t stream_cb{0}; diff --git a/families/personaplex/system_prompt.py b/families/personaplex/system_prompt.py new file mode 100644 index 0000000000..66944384e9 --- /dev/null +++ b/families/personaplex/system_prompt.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Prepare the forced PersonaPlex voice/text prefix without running a language model. + +The official .pt voice contains summed input embeddings and its delayed token +cache. Subsequent system-prompt steps force every stream, so their input +embeddings and final cache do not depend on sampled model outputs. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +DELAYS = np.array([0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1]) +SILENCE = np.array([948, 243, 1178, 546, 1736, 1030, 1978, 2008]) +SINE = np.array([430, 1268, 381, 1611, 1095, 1495, 56, 472]) + + +def wrap_system_text(text: str) -> str: + text = text.strip() + return text if text.startswith("") and text.endswith("") else f" {text} " + + +class ForcedPromptCache: + """The official LMGen delay-cache transition when all streams are forced.""" + + def __init__(self, text_bos: int, audio_bos: int = 2048): + self.initial = np.array([text_bos] + [audio_bos] * 16) + self.cache = np.full((17, 4), -2, dtype=np.int64) + self.provided = np.zeros((17, 4), dtype=bool) + self.offset = 0 + + def step(self, user: np.ndarray, agent: np.ndarray, text: int) -> np.ndarray | None: + values = np.concatenate(([text], agent, user)) + for k, delay in enumerate(DELAYS): + pos = (self.offset + delay) % 4 + self.cache[k, pos] = values[k] + self.provided[k, pos] = True + if self.offset <= delay: + self.cache[k, self.offset % 4] = self.initial[k] + self.provided[k, self.offset % 4] = True + if self.offset == 0: + self.cache[:, 0] = self.initial + self.offset = 1 + return None + source, target = (self.offset - 1) % 4, self.offset % 4 + inputs = self.cache[:, source].copy() + if not self.provided[:, target].all() or (inputs == -2).any(): + raise ValueError("a system-prompt step requires unforced model predictions") + self.provided[:, source] = False + self.offset += 1 + return inputs + + +def prepare(checkpoint: Path, voice: Path, text: str, precision: str, output: Path) -> dict: + import sentencepiece + import torch + from safetensors import safe_open + + dtype = {"fp16": torch.float16, "fp32": torch.float32}[precision] + tokenizer = sentencepiece.SentencePieceProcessor(str(checkpoint / "tokenizer_spm_32k_3.model")) + text_ids = tokenizer.encode(wrap_system_text(text)) + saved = torch.load(voice, map_location="cpu", weights_only=True) + embeddings = saved["embeddings"].to(dtype).float().reshape(-1, 4096).numpy() + voice_cache = saved["cache"].cpu().numpy() + if voice_cache.shape != (1, 17, 4) or not np.isfinite(embeddings).all(): + raise ValueError("unsupported official voice-prompt shape or non-finite embeddings") + state = ForcedPromptCache(32000) + bos = np.full(8, 2048) + # step_embeddings retries the offset-zero preparation before its first forward. + for _ in embeddings: + while state.step(bos, bos, 3) is None: + pass + state.cache[:] = voice_cache[0] + prefix = [row for row in embeddings] + with safe_open(checkpoint / "model.safetensors", framework="pt", device="cpu") as weights: + rows: dict[tuple[int, int], torch.Tensor] = {} + + def embed(tokens: np.ndarray) -> np.ndarray: + result = None + # Native embed_codes sums the audio tables first, then the text table. + for k in [*range(1, 17), 0]: + token = int(tokens[k]) + key = k, token + if key not in rows: + name = "text_emb.weight" if k == 0 else f"emb.{k - 1}.weight" + row = weights.get_slice(name)[max(token, 0)].to(dtype) + rows[key] = torch.zeros_like(row) if token == -1 else row + result = rows[key] if result is None else result + rows[key] + return result.float().numpy() + + for token in [*[3] * 6, *text_ids, *[3] * 6]: + inputs = state.step(SINE, SILENCE, token) + if inputs is None: + raise ValueError("unexpected uninitialized voice-prefix state") + prefix.append(embed(inputs)) + metadata = {"version": 1, "precision": precision, "voice": voice.name, + "text_prompt": text, "text_token_ids": text_ids, "voice_frames": len(embeddings), + "offset": state.offset, "frames": len(prefix), "silence_frames": 6} + output.parent.mkdir(parents=True, exist_ok=True) + np.savez(output, embeddings=np.asarray(prefix, dtype=np.float32), cache=state.cache.astype(np.int32), + provided=state.provided.astype(np.uint8), metadata=np.array(json.dumps(metadata))) + return metadata + + +def load(path: Path, precision: str, hidden: int, capacity: int) -> tuple[dict, bytes]: + """Validate one prepared prefix before building any engine.""" + with np.load(path, allow_pickle=False) as data: + metadata = json.loads(str(data["metadata"])) + embeddings = np.asarray(data["embeddings"], dtype=np.float32) + cache, provided = data["cache"], data["provided"] + if (metadata["version"] != 1 or metadata["precision"] != precision + or embeddings.shape != (metadata["frames"], hidden) + or metadata["offset"] != metadata["frames"] + 1 + or metadata["offset"] >= capacity or not np.isfinite(embeddings).all() + or cache.shape != (17, 4) or provided.shape != cache.shape + or not np.isin(provided, [0, 1]).all()): + raise ValueError("system prompt does not match the engine precision, geometry or capacity") + return {**metadata, "cache": cache.reshape(-1).tolist(), + "provided": provided.reshape(-1).tolist()}, embeddings.tobytes() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--voice-prompt", type=Path, required=True) + parser.add_argument("--text-prompt", required=True) + parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + print(json.dumps(prepare(args.checkpoint, args.voice_prompt, args.text_prompt, args.precision, args.output))) diff --git a/families/personaplex/tests/aiperf/README.md b/families/personaplex/tests/aiperf/README.md new file mode 100644 index 0000000000..8cfffe40e1 --- /dev/null +++ b/families/personaplex/tests/aiperf/README.md @@ -0,0 +1,73 @@ +# PersonaPlex speech question answering + +The `personaplex-7b` AIPerf profile uses all 455 questions from VoiceBench +OpenBookQA, pinned to dataset revision +`b02edcef1330480be3a11bd6f7434ac32f05ad08`. Each question has an objective +A/B/C/D answer. Only the question audio enters the model; the question text and +gold answer are available to the scorer only. The suite resamples the 16 kHz +questions once to Mimi's 24 kHz rate (SciPy polyphase), so both sides read the +same samples and neither pipeline resamples. + +Both sides use the official NATM1 voice and the same teacher role prompt, +greedy decoding, an FP16 temporal/depth model, and an FP32 Mimi codec. The +official voice asset's BF16 embeddings are cast to FP16 on both sides. Each +request resets its state and processes the full question followed by 125 +silent input frames (10 seconds) for the reply. The prepared bundle's temporal +capacity covers the longest question, reply window and system prompt. Original +weights remain in the shared HF cache; the prepared model directory contains +symlinks, one config overlay and small derived prompt files. + +Native applies the pinned official `step_system_prompts` implementation. TRTMC +replays the same input embeddings through its temporal engine, restores the +official delayed-token state, and then processes the user stream. Every stream +is forced during the prefix, so its embeddings and token cache can be prepared +without generating predictions. Prefix execution is included in each Task +call; preparation, model loading and output scoring are excluded. + +During the question, both implementations force the agent's text PAD and official +silence audio tokens while processing the real user audio. Generation is released +for the reply window, with the EPAD text token (0) triggering the response as in +[Moshi's spoken-QA protocol](https://arxiv.org/html/2410.00037v2#S5.SS5). +This gives the model the complete question and all options +before answering, as required by this turn-based QA protocol. Original full-duplex +barge-in behavior remains available when `listen_before_answer` is disabled. + +The same AIPerf inference outputs provide both results: + +- **Acc:** transcribe the generated audio using pinned + `openai/whisper-large-v3-turbo@41f01f3fe87f28c78e2fbf8b568835947dd65ed9`, + evaluating the final 10-second answer window while retaining complete raw audio, + then match an explicit option letter or a unique complete option text. + Explicit answer declarations take precedence over words in explanations; + conflicting declarations and negated option mentions do not become answers. + Empty, ambiguous and unparsable replies count as wrong. The ASR process + receives output paths only, without gold labels. Native and TRTMC each get + an absolute score and a paired comparison with a 2-point regression margin; + there is no minimum Native score. ASR errors remain a measurement limitation. +- **Perf:** Native and TRTMC Task-call P50 over those same questions, after + excluded warmup. Input/reply lengths and output frame counts must agree. + Offline ASR and grading do not enter either P50. AIPerf client measurements + remain informational. + +Run the standard profile: + +```bash +trtmc-aiperf-qual run --profile personaplex-7b \ + --environment environment.yaml --out results/personaplex-7b +``` + +For a pilot, resolve the profile and set +`absolute[0].suite_definition.selection.count` to 50 before `campaign.run_one`. +Keep the pilot and full results in separate directories. Preserve generated +audio for scoring and review with `retention.temporary_files: retain`; bundles +can be deleted after qualification while keeping `retention.hf_cache: retain`. + +This evaluates spoken single-question answering. It does not establish +interruption handling, turn-taking quality or voice similarity. Previous +LibriSpeech waveform-parity results remain separate historical diagnostics; +their scores and passing criteria are not reinterpreted as QA accuracy. + +The profile's own requirements file pins dataset preparation and ASR dependencies +in its reference environment. Offline postprocessing uses that interpreter by +default; `scoring_python` in the environment can select a separate compatible +scoring interpreter. Neither path changes the measured inference duration. diff --git a/families/personaplex/tests/aiperf/personaplex-7b.yaml b/families/personaplex/tests/aiperf/personaplex-7b.yaml new file mode 100644 index 0000000000..a3c4ff6302 --- /dev/null +++ b/families/personaplex/tests/aiperf/personaplex-7b.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +candidate: + model_directory: trtmc-reference/personaplex-model + build: {fp32_layers: []} +reference: + precision: fp16 + requirements: families/personaplex/tests/requirements.txt + prepare: families/personaplex/reference/prepare_voicebench.py + options: + listen_before_answer: true + system_prompt: trtmc-reference/personaplex-model/system_prompt.npz + voice_prompt: trtmc-reference/personaplex-model/NATM1.pt +absolute: [voicebench_speech_qa] +benchmarks: + voicebench_speech_qa: + suite: voicebench-openbookqa + suite_definition: + suite: voicebench-openbookqa-455 + version: 1 + source: + kind: hf_dataset + dataset: hlt-lab/voicebench + subset: openbookqa + split: test + revision: b02edcef1330480be3a11bd6f7434ac32f05ad08 + audio_column: audio + audio_sample_rate: 24000 # Mimi's rate: both sides read the same resampled 16 kHz questions + label_fields: {answer: reference, prompt: prompt} + selection: {method: first, count: 455} + fields: {audio: audio_path} + request: {max_new_tokens: 1024, tail_frames: 125} + metric: spoken_choice + gate: {margin: 2.0} + postprocess: + script: apps/aiperf_qual/trtmc_aiperf_qual/transcribe_outputs.py + arguments: [--model, openai/whisper-large-v3-turbo, --revision, 41f01f3fe87f28c78e2fbf8b568835947dd65ed9, + --last-seconds, "10"] diff --git a/families/personaplex/tests/cpp/test_speech_generation_helpers.cpp b/families/personaplex/tests/cpp/test_speech_generation_helpers.cpp index 39ab651c2c..631c4d49a7 100644 --- a/families/personaplex/tests/cpp/test_speech_generation_helpers.cpp +++ b/families/personaplex/tests/cpp/test_speech_generation_helpers.cpp @@ -111,6 +111,32 @@ void test_waveform_trim_and_peak_normalize() { check_close(waveform[1], -0.475F, 1e-6F, "waveform normalize sample 1"); } +void test_prefilled_cache_writes_the_first_user_frame_at_the_absolute_offset() { + auto state = trtmc::make_delay_cache_state({}, 16); + const std::vector codec_tokens{101, 102, 103, 104, 105, 106, 107, 108}; + trtmc::write_user_tokens_to_delay_cache(state, codec_tokens, 109, 8, 1, 8, 2048, 0); + for (int32_t cb = 0; cb < 8; ++cb) { + const int32_t stream = 9 + cb; + const auto index = trtmc::delay_cache_index(state, stream, 109 + state.delays[stream]); + check(state.cache[index] == codec_tokens[cb], "prefilled cache reads relative user frame"); + check(state.provided[index] == 1, "prefilled cache forces the user stream"); + } +} + +void test_listening_forces_only_the_agent_streams_at_their_delays() { + auto state = trtmc::make_delay_cache_state({}, 16); + trtmc::write_user_tokens_to_delay_cache(state, std::vector(8, 101), 109, 8, 1, 8, 2048, + 0); + trtmc::write_listening_tokens_to_delay_cache(state, 109, 3); + const std::vector agent{3, 948, 243, 1178, 546, 1736, 1030, 1978, 2008}; + for (int32_t k = 0; k < 17; ++k) { + const auto index = trtmc::delay_cache_index(state, k, 109 + state.delays[k]); + check(state.cache[index] == (k < 9 ? agent[k] : 101), + "listening retains user audio and forces official agent silence"); + check(state.provided[index] == 1, "listening marks each forced stream"); + } +} + void test_waveform_postprocess_skips_invalid_or_safe_inputs() { std::vector waveform = {0.1F, 0.2F, 0.3F}; std::vector empty_waveform; @@ -200,6 +226,8 @@ void test_mimi_ring_attention_matches_official_cache_wrap() { int main() { test_default_speech_delays_generalize_to_num_codebooks(); test_delay_cache_reads_and_collects_outputs(); + test_prefilled_cache_writes_the_first_user_frame_at_the_absolute_offset(); + test_listening_forces_only_the_agent_streams_at_their_delays(); test_waveform_trim_and_peak_normalize(); test_waveform_postprocess_skips_invalid_or_safe_inputs(); test_mimi_encode_plan_keeps_only_the_causal_input_prefix(); diff --git a/families/personaplex/tests/requirements.txt b/families/personaplex/tests/requirements.txt new file mode 100644 index 0000000000..7c3c846eca --- /dev/null +++ b/families/personaplex/tests/requirements.txt @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Native pipeline, dataset preparation and offline scoring in one isolated environment. +einops==0.7.0 +soundfile==0.14.0 +datasets==5.0.0 +pyarrow==24.0.0 +librosa==0.11.0 diff --git a/families/personaplex/tests/test_reference_decode_work.py b/families/personaplex/tests/test_reference_decode_work.py new file mode 100644 index 0000000000..2f104d5de0 --- /dev/null +++ b/families/personaplex/tests/test_reference_decode_work.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The timed Native task decodes the audio it returns, once per frame, from unresampled input.""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace +import wave + +import numpy as np +import pytest + +from families.personaplex.reference.adapter import Adapter + + +def _write_wav(path, rate, samples): + with wave.open(str(path), "wb") as stream: + stream.setnchannels(1) + stream.setsampwidth(2) + stream.setframerate(rate) + stream.writeframes(np.zeros(samples, dtype=" Teacher. " + assert wrap_system_text(" Teacher. ") == " Teacher. " + + +def test_delayed_streams_read_the_preceding_frame(): + state = ForcedPromptCache(32000) + assert state.step(SINE, SILENCE, 3) is None + first = state.step(SINE, SILENCE, 42) + np.testing.assert_array_equal(first, [32000] + [2048] * 16) + second = state.step(SINE, SILENCE, 43) + assert second[0] == 42 + assert second[1] == SILENCE[0] + assert second[2] == 2048 # one-step-delayed agent stream + third = state.step(SINE, SILENCE, 44) + np.testing.assert_array_equal(third[1:9], SILENCE) + np.testing.assert_array_equal(third[9:], SINE) + assert third[0] == 43 + + +@pytest.mark.parametrize("precision,capacity", [("fp32", 512), ("fp16", 3)]) +def test_prefix_rejects_another_dtype_or_insufficient_context(tmp_path, precision, capacity): + p = tmp_path / "prompt.npz" + np.savez(p, embeddings=np.zeros((2, 4), dtype=np.float32), cache=np.zeros((17, 4), dtype=np.int32), + provided=np.zeros((17, 4), dtype=np.uint8), + metadata=np.array(json.dumps({"version": 1, "precision": "fp16", "frames": 2, "offset": 3}))) + with pytest.raises(ValueError, match="does not match"): + load(p, precision, 4, capacity) + + +@pytest.mark.parametrize("listen", [False, True]) +def test_prompt_and_full_reply_window_are_inside_the_native_task(monkeypatch, tmp_path, listen): + torch = pytest.importorskip("torch") + soundfile = pytest.importorskip("soundfile") + from families.personaplex.reference.adapter import Adapter + + events = [] + steps = [] + + class Codec: + sample_rate = 24_000 + + def reset_streaming(self): + events.append("codec-reset") + + def decode(self, codes): + return torch.zeros((1, 1, 1920)) + + class Generator: + zero_text_code = 3 + + def _encode_zero_frame(self): + return torch.tensor(SILENCE).reshape(1, 8, 1) + + def reset_streaming(self): + events.append("generator-reset") + + def step_system_prompts(self, codec): + events.append("prompt") + + def step(self, codes, **forced): + steps.append(forced) + return torch.zeros((1, 17, 1), dtype=torch.int64) + + def iterate(audio, **kwargs): + assert audio.shape == (1, 3 * 1920) + assert np.count_nonzero(audio[:, 1920:]) == 0 + return range(3) + + lm = ModuleType("moshi.models.lm") + lm._iterate_audio = iterate + lm.encode_from_sphn = lambda codec, frames, **kwargs: (torch.zeros((1, 8, 1)) for _ in frames) + monkeypatch.setitem(sys.modules, "moshi.models.lm", lm) + + def timed(call): + events.append("timing-start") + result = call() + events.append("timing-end") + return result, 1.0 + + adapter = Adapter.__new__(Adapter) + adapter.frame_size = 1920 + adapter.mimi, adapter.other_mimi, adapter.generator = Codec(), Codec(), Generator() + adapter.system_prompt = {"voice": "NATM1.pt", "frames": 108} + adapter.listen_before_answer = listen + adapter.host = SimpleNamespace(required=lambda r, k: r[k], load_audio=lambda *a: np.ones(1920), + timed=timed, tensor_observation=lambda audio, base: {"samples": audio.size}, + invocation=lambda observation, ms: observation) + question = tmp_path / "question.wav" # the reference only checks its rate; load_audio is stubbed + soundfile.write(question, np.zeros(1920, dtype=np.float32), 24_000) + result = adapter.invoke({"audio_path": str(question), "tail_frames": 2, "max_new_tokens": 10}, tmp_path / "out") + assert result["output_tokens"] == 3 and result["samples"] == 3 * 1920 + assert result["system_prompt_voice"] == "NATM1.pt" + assert bool(steps[0]) is listen + assert steps[1:] == ([{"text_token": 0}, {}] if listen else [{}, {}]) + if listen: + assert steps[0]["text_token"] == 3 + np.testing.assert_array_equal(steps[0]["moshi_tokens"].numpy().reshape(-1), SILENCE) + assert events == ["timing-start", "codec-reset", "codec-reset", "generator-reset", "prompt", + "codec-reset", "timing-end"]