Spaces:
Running on Zero
Running on Zero
| 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,)) | |
| 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() | |