wehub-resource-sync
ba4be087d5
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
CICD NeMo / pre-flight (push) Failing after 0s
CICD NeMo / configure (push) Has been skipped
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
CICD NeMo / code-linting (push) Has been skipped
Build, validate, and release Neural Modules / release (push) Has been skipped
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
CICD NeMo / cicd-test-container-build (push) Has been cancelled
CICD NeMo / cicd-import-tests (push) Has been cancelled
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
CICD NeMo / cicd-main-speech (push) Has been cancelled
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
CICD NeMo / Coverage (e2e) (push) Has been cancelled
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
86 lines
2.6 KiB
Python
86 lines
2.6 KiB
Python
# Copyright (c) 2025, NVIDIA CORPORATION.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Functional tests for titanet_large."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
MODEL_NAME = "titanet_large"
|
|
NEMO_FILE = "titanet_large.nemo"
|
|
|
|
MODEL_DIR = os.environ.get(
|
|
"NEMO_MODEL_SUPPORT_DIR",
|
|
os.environ.get("NEMO_MODEL_SUPPORT_DIR_CI", "/home/TestData/nemo-speech-ci-models"),
|
|
)
|
|
_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
_model = None
|
|
|
|
|
|
def _load_model():
|
|
global _model
|
|
if _model is not None:
|
|
return _model
|
|
from nemo.collections.asr.models import EncDecSpeakerLabelModel
|
|
|
|
filepath = os.path.join(MODEL_DIR, NEMO_FILE)
|
|
_model = EncDecSpeakerLabelModel.restore_from(filepath, map_location="cpu").to(_DEVICE)
|
|
return _model
|
|
|
|
|
|
def test_model_init():
|
|
model = _load_model()
|
|
assert model is not None
|
|
if hasattr(model, "to_config_dict"):
|
|
cfg = model.to_config_dict()
|
|
assert cfg is not None
|
|
|
|
|
|
def test_model_training_step():
|
|
"""Run one training step via direct training_step() call."""
|
|
from e2e_utils import prepare_for_training_step
|
|
|
|
model = _load_model()
|
|
prepare_for_training_step(model)
|
|
d = next(model.parameters()).device
|
|
num_classes = model.decoder._num_classes
|
|
batch = (
|
|
torch.randn(2, 16000, device=d),
|
|
torch.tensor([16000, 12000], device=d),
|
|
torch.randint(0, num_classes, (2,), device=d),
|
|
torch.tensor([1, 1], dtype=torch.long, device=d),
|
|
)
|
|
result = model.training_step(batch, 0)
|
|
loss = result if isinstance(result, torch.Tensor) else result['loss']
|
|
assert torch.isfinite(loss), f"Loss is not finite: {loss}"
|
|
loss.backward()
|
|
|
|
|
|
def test_model_inference():
|
|
model = _load_model()
|
|
model.eval()
|
|
d = _DEVICE
|
|
with torch.no_grad():
|
|
logits, embs = model.forward(
|
|
input_signal=torch.randn(1, 16000, device=d),
|
|
input_signal_length=torch.tensor([16000], device=d),
|
|
)
|
|
assert logits is not None
|
|
assert embs is not None
|
|
assert embs.ndim == 2 # (B, embedding_dim)
|
|
assert embs.shape[0] == 1
|
|
assert torch.isfinite(embs).all()
|