Spaces:
Running on Zero
Running on Zero
File size: 5,509 Bytes
82c62b3 2049495 82c62b3 2049495 | 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 | import shutil
import subprocess
import tempfile
import unittest
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from fireredtts3.reference_asr import (
QWEN3_ASR_MODEL_ID,
transcribe_audio,
transcribe_reference_audio,
)
REFERENCE_TRANSCRIPT = "The reference transcript is filled automatically."
class _FakeBatch(dict):
def to(self, device, dtype):
self.to_args = (device, dtype)
return self
class _FakeProcessor:
def __init__(self, decoded_text=REFERENCE_TRANSCRIPT):
self.decoded_text = decoded_text
self.audio = None
self.sampling_rate = None
self.decode_args = None
def apply_transcription_request(self, audio, processor_kwargs):
self.audio = audio
self.sampling_rate = processor_kwargs["sampling_rate"]
return _FakeBatch({"input_ids": torch.tensor([[10, 11, 12]])})
def decode(self, generated_ids, return_format):
self.decode_args = (generated_ids, return_format)
return [self.decoded_text]
class _FakeModel:
device = torch.device("cpu")
dtype = torch.float32
def __init__(self):
self.generate_kwargs = None
def generate(self, **kwargs):
self.generate_kwargs = kwargs
return torch.tensor([[10, 11, 12, 20, 21]])
class TranscribeAudioUnitTests(unittest.TestCase):
def test_returns_only_the_stripped_reference_transcript(self):
processor = _FakeProcessor(f" {REFERENCE_TRANSCRIPT} ")
model = _FakeModel()
transcript = transcribe_audio(
np.zeros(16000, dtype=np.float32),
16000,
processor=processor,
model=model,
)
self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
self.assertEqual(processor.decode_args[1], "transcription_only")
self.assertEqual(processor.decode_args[0].tolist(), [[20, 21]])
self.assertEqual(model.generate_kwargs["max_new_tokens"], 256)
self.assertFalse(model.generate_kwargs["do_sample"])
def test_rejects_an_empty_model_transcript(self):
with self.assertRaisesRegex(RuntimeError, "empty transcript"):
transcribe_audio(
np.zeros(16000, dtype=np.float32),
16000,
processor=_FakeProcessor(" "),
model=_FakeModel(),
)
def test_rejects_invalid_audio_input(self):
with self.assertRaisesRegex(ValueError, "mono waveform"):
transcribe_audio(
np.zeros((1, 16000), dtype=np.float32),
16000,
processor=_FakeProcessor(),
model=_FakeModel(),
)
class ReferenceTranscriptIntegrationTests(unittest.TestCase):
def test_uploaded_audio_is_transcribed_into_reference_transcript(self):
processor = _FakeProcessor()
model = _FakeModel()
with tempfile.TemporaryDirectory() as tmpdir:
audio_path = Path(tmpdir) / "reference.wav"
sf.write(audio_path, np.zeros(8000, dtype=np.float32), 8000)
transcript = transcribe_reference_audio(
audio_path,
processor=processor,
model=model,
)
self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
self.assertEqual(processor.sampling_rate, 16000)
self.assertEqual(processor.audio.shape, (16000,))
def test_reference_audio_is_limited_to_the_clone_prompt_duration(self):
processor = _FakeProcessor()
with tempfile.TemporaryDirectory() as tmpdir:
audio_path = Path(tmpdir) / "long-reference.wav"
sf.write(audio_path, np.zeros(25 * 8000, dtype=np.float32), 8000)
transcribe_reference_audio(
audio_path,
processor=processor,
model=_FakeModel(),
max_seconds=20.0,
)
self.assertEqual(processor.audio.shape, (20 * 16000,))
@unittest.skipUnless(shutil.which("ffmpeg"), "FFmpeg is required for M4A coverage")
def test_m4a_upload_is_transcribed_into_reference_transcript(self):
processor = _FakeProcessor()
model = _FakeModel()
with tempfile.TemporaryDirectory() as tmpdir:
wav_path = Path(tmpdir) / "reference.wav"
m4a_path = Path(tmpdir) / "reference clip.m4a"
time = np.arange(8000, dtype=np.float32) / 8000
sf.write(wav_path, 0.25 * np.sin(2 * np.pi * 440 * time), 8000)
subprocess.run(
[
"ffmpeg",
"-nostdin",
"-loglevel",
"error",
"-y",
"-i",
str(wav_path),
"-c:a",
"aac",
str(m4a_path),
],
check=True,
)
transcript = transcribe_reference_audio(
m4a_path,
processor=processor,
model=model,
)
self.assertEqual(transcript, REFERENCE_TRANSCRIPT)
self.assertEqual(processor.sampling_rate, 16000)
self.assertGreaterEqual(processor.audio.shape[0], 16000)
self.assertLess(processor.audio.shape[0], 17000)
def test_uses_the_official_hugging_face_qwen3_asr_checkpoint(self):
self.assertEqual(QWEN3_ASR_MODEL_ID, "Qwen/Qwen3-ASR-0.6B-hf")
if __name__ == "__main__":
unittest.main()
|