File size: 7,456 Bytes
4532b62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | # ---------------------------------------------------------------------
# Copyright (c) 2026 Qualcomm Technologies, Inc. and/or its subsidiaries.
# SPDX-License-Identifier: BSD-3-Clause
# ---------------------------------------------------------------------
from __future__ import annotations
from pathlib import Path
import numpy as np
import pytest
import soundfile as sf
import torch
from huggingface_hub.errors import GatedRepoError
from .demo import main as demo_main
from .model import (
CONTEXT_LENGTH,
DECODE_SEQ_LEN,
HEAD_DIM,
NUM_KEY_VALUE_HEADS,
NUM_LAYERS,
PREFILL_SEQ_LEN,
SAMPLE_RATE,
NeuTTSNano,
build_attention_mask,
empty_kv_cache,
)
# The NeuTTS-Nano backbone is a gated HF repo. Skip rather than fail on machines
# whose token has not accepted the NeuTTS Open License.
_GATED_SKIP = "NeuTTS-Nano weights are HF-gated for this token"
def _load_or_skip() -> NeuTTSNano:
try:
return NeuTTSNano.from_pretrained()
except GatedRepoError:
pytest.skip(_GATED_SKIP)
def _run(model: NeuTTSNano, seq_len: int) -> list[torch.Tensor]:
ids = torch.zeros((1, seq_len), dtype=torch.int32)
position_ids = torch.arange(seq_len).reshape(1, seq_len)
cos, sin = model.backbone.embedding.get_embedding(position_ids)
caches = empty_kv_cache(CONTEXT_LENGTH - seq_len)
with torch.no_grad():
return model.backbone(
ids, build_attention_mask(seq_len, seq_len), cos, sin, *caches
)
@pytest.mark.slow
@pytest.mark.parametrize("seq_len", [PREFILL_SEQ_LEN, DECODE_SEQ_LEN])
def test_backbone_graph_forward(seq_len: int) -> None:
model = _load_or_skip()
out = _run(model, seq_len)
vocab_size = model.backbone.llm_config.vocab_size
assert out[0].shape == (1, seq_len, vocab_size)
assert torch.isfinite(out[0]).all()
# logits plus one key and one value per layer
assert len(out) == 1 + 2 * NUM_LAYERS
for layer in range(NUM_LAYERS):
key, value = out[1 + 2 * layer], out[2 + 2 * layer]
assert key.shape == (NUM_KEY_VALUE_HEADS, 1, HEAD_DIM, seq_len)
assert value.shape == (NUM_KEY_VALUE_HEADS, 1, seq_len, HEAD_DIM)
assert torch.isfinite(key).all() and torch.isfinite(value).all()
@pytest.mark.slow
def test_graphs_share_one_source() -> None:
"""Both graphs must describe the same weights at two sequence lengths."""
model = _load_or_skip()
assert model.graph_names == [model.prefill_graph, model.decode_graph]
assert model.shared_source_model
prefill = model.get_graph_input_spec(model.prefill_graph)
decode = model.get_graph_input_spec(model.decode_graph)
# Same inputs, same cache contract; only the sequence dimension differs.
assert set(prefill) == set(decode)
assert prefill["input_ids"][0] == (1, PREFILL_SEQ_LEN)
assert decode["input_ids"][0] == (1, DECODE_SEQ_LEN)
assert model.get_graph_output_spec(model.prefill_graph) == (
model.get_graph_output_spec(model.decode_graph)
)
@pytest.mark.slow
def test_shared_source_serializes_once(tmp_path: Path) -> None:
"""One .pt must serve both graphs, and be traced at the longest shape."""
model = _load_or_skip()
path = model.serialize_graph(model.decode_graph, tmp_path)
assert path.is_file()
traced = torch.jit.load(str(path))
inputs = [
torch.from_numpy(v[0])
for v in model.get_graph_sample_inputs(model.prefill_graph).values()
]
with torch.no_grad():
out = traced(*inputs)
assert out[0].shape[1] == PREFILL_SEQ_LEN
@pytest.mark.slow
def test_decode_matches_prefill() -> None:
"""Cached decode of the last token must match prefilling the whole window.
Predicting what follows tokens ``0..n-1`` is done two ways: one prefill over
all ``n``, versus a prefill over ``0..n-2`` (left-padded) followed by a
single-token decode. Agreement exercises the cache layout, the sliding
window and the mask together.
"""
model = _load_or_skip()
backbone = model.backbone
rng = np.random.default_rng(seed=0)
ids = torch.from_numpy(
rng.integers(low=0, high=1000, size=(1, PREFILL_SEQ_LEN)).astype(np.int32)
)
split = PREFILL_SEQ_LEN - 1
cos, sin = backbone.embedding.get_embedding(
torch.arange(PREFILL_SEQ_LEN).reshape(1, -1)
)
with torch.no_grad():
whole = backbone(
ids,
build_attention_mask(PREFILL_SEQ_LEN, PREFILL_SEQ_LEN),
cos,
sin,
*empty_kv_cache(CONTEXT_LENGTH - PREFILL_SEQ_LEN),
)
# Left-pad by one so the first ``split`` real tokens stay right-aligned.
padded = torch.cat([torch.zeros((1, 1), dtype=torch.int32), ids[:, :split]], dim=1)
cos, sin = backbone.embedding.get_embedding(torch.tensor([[0, *range(split)]]))
with torch.no_grad():
partial = backbone(
padded,
build_attention_mask(PREFILL_SEQ_LEN, split),
cos,
sin,
*empty_kv_cache(CONTEXT_LENGTH - PREFILL_SEQ_LEN),
)
# Drop the pad position's cache entry, then right-align in the decode buffer.
caches = empty_kv_cache(CONTEXT_LENGTH - 1)
for layer in range(NUM_LAYERS):
caches[2 * layer][..., -split:] = partial[1 + 2 * layer][..., 1:]
caches[2 * layer + 1][:, :, -split:, :] = partial[2 + 2 * layer][:, :, 1:, :]
cos, sin = backbone.embedding.get_embedding(torch.tensor([[split]]))
with torch.no_grad():
step = backbone(
ids[:, split:],
build_attention_mask(1, PREFILL_SEQ_LEN),
cos,
sin,
*caches,
)
torch.testing.assert_close(step[0][:, 0], whole[0][:, split], atol=1e-3, rtol=1e-3)
def _spectral_centroid_std(wav: np.ndarray, sample_rate: int) -> float:
"""Standard deviation of the spectral centroid over time, in Hz.
Separates speech from a stationary drone. librosa would do this in one call
but imports numba, so compute it with numpy.
"""
frame, hop = 1024, 256
n = 1 + max(0, (len(wav) - frame) // hop)
idx = np.arange(frame)[None, :] + hop * np.arange(n)[:, None]
mag = np.abs(np.fft.rfft(wav[idx] * np.hanning(frame), axis=1))
freqs = np.fft.rfftfreq(frame, 1.0 / sample_rate)
centroid = (mag * freqs).sum(axis=1) / (mag.sum(axis=1) + 1e-10)
return float(centroid.std())
@pytest.mark.slow
def test_demo(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
# The demo fixes its own sampling seed, so this is deterministic. Seeding from
# here would not work: upstream's infer() calls torch.manual_seed itself.
monkeypatch.chdir(tmp_path)
try:
demo_main(is_test=True)
except GatedRepoError:
pytest.skip(_GATED_SKIP)
output = tmp_path / "neutts_nano_output.wav"
assert output.is_file()
wav, sample_rate = sf.read(output, dtype="float32", always_2d=True)
wav = wav.mean(axis=1)
assert sample_rate == SAMPLE_RATE
assert np.isfinite(wav).all()
# Guard against a silent / degenerate waveform passing as success.
assert len(wav) > SAMPLE_RATE
assert np.abs(wav).max() > 1e-2
# A broken cache or mask collapses generation into a short repeating token
# loop, which decodes to a stationary drone: real speech measures >1300 Hz of
# centroid spread, a greedy-collapsed loop 47 Hz.
assert _spectral_centroid_std(wav, sample_rate) > 300.0
|