diff --git a/README.md b/README.md index 9ff28cd..5f606a4 100644 --- a/README.md +++ b/README.md @@ -236,6 +236,14 @@ The message flow follows the common Qwen multimodal pattern. The chat template i 3. `processor(text=text, audio=audios)` computes Whisper input features and expands audio placeholders. 4. `model.generate(...)` produces timestamped transcription and diarization text. +`build_transcription_messages` also accepts a one-dimensional `numpy.ndarray` or `torch.Tensor` waveform, in addition to a local path. The waveform should be mono floating-point samples at the processor's configured sampling rate: + +```python +messages = build_transcription_messages(waveform, prompt="Transcribe and diarize this audio.") +``` + +The helper preserves the array in the message and converts tensors to NumPy before handing them to the Transformers audio loader. + The repository's Hugging Face loader uses an explicit attention policy instead of accepting a silent Transformers fallback: `flash_attention_4`, `flash_attention_3`, `flash_attention_2`, `sdpa`, then `eager`. Unavailable candidates and the selected diff --git a/README_zh.md b/README_zh.md index 30b2a9a..3310e0e 100644 --- a/README_zh.md +++ b/README_zh.md @@ -236,6 +236,14 @@ for segment in parse_transcript(result["text"]): 3. `processor(text=text, audio=audios)` 计算 Whisper 输入特征并展开音频占位符。 4. `model.generate(...)` 生成带时间戳的转写与说话人分离文本。 +除本地文件路径外，`build_transcription_messages` 也接受一维 `numpy.ndarray` 或 `torch.Tensor` 波形。波形应为单声道浮点采样，并使用 processor 配置的采样率： + +```python +messages = build_transcription_messages(waveform, prompt="Transcribe and diarize this audio.") +``` + +该辅助函数会保留消息中的数组，并在交给 Transformers 音频 loader 前将 Tensor 转换为 NumPy。 + 仓库的 Hugging Face 加载器使用显式的 attention 策略，不接受 Transformers 的静默回退：候选顺序为 `flash_attention_4`、`flash_attention_3`、`flash_attention_2`、`sdpa`，最后才是 `eager`。 不可用的候选和最终选择都会写入日志；如果最终落到 `eager`，会额外发出 warning，因为它会为长音频物化二次方大小的 attention 张量。 diff --git a/moss_transcribe_diarize/inference_utils.py b/moss_transcribe_diarize/inference_utils.py index 81d8164..fd94f82 100644 --- a/moss_transcribe_diarize/inference_utils.py +++ b/moss_transcribe_diarize/inference_utils.py @@ -112,8 +112,10 @@ def load_audio_av(audio: str, sampling_rate: int) -> np.ndarray: return (np.concatenate(chunks).astype(np.float32) / 32768.0).astype(np.float32, copy=False) -def load_audio_item(audio: str | np.ndarray, sampling_rate: int) -> np.ndarray: +def load_audio_item(audio: str | np.ndarray | torch.Tensor, sampling_rate: int) -> np.ndarray: """Load audio with Transformers' loader, using PyAV for media containers.""" + if torch.is_tensor(audio): + audio = audio.detach().cpu().numpy() if isinstance(audio, str) and _is_likely_video_path(audio): return load_audio_av(audio, sampling_rate=sampling_rate) try: @@ -139,19 +141,31 @@ def process_audio_info(messages: list[dict[str, Any]], sampling_rate: int): for item in content: if item.get("type") != "audio": continue - audio = item.get("audio") or item.get("audio_url") or item.get("url") or item.get("path") + # Do not use ``or`` here: NumPy arrays and tensors have no single + # truth value when they contain more than one sample. + audio = None + for key in ("audio", "audio_url", "url", "path"): + value = item.get(key) + if value is not None: + audio = value + break if audio is None: raise ValueError("Audio content must include audio, audio_url, url, or path.") audios.append(load_audio_item(audio, sampling_rate=sampling_rate)) return audios -def build_transcription_messages(audio_path: str | Path, prompt: str = DEFAULT_PROMPT) -> list[dict[str, Any]]: +def build_transcription_messages( + audio_path: str | Path | np.ndarray | torch.Tensor, + prompt: str = DEFAULT_PROMPT, +) -> list[dict[str, Any]]: + """Build a chat message for a path or an in-memory mono waveform.""" + audio = str(audio_path) if isinstance(audio_path, Path) else audio_path return [ { "role": "user", "content": [ - {"type": "audio", "audio": str(audio_path)}, + {"type": "audio", "audio": audio}, {"type": "text", "text": prompt.strip() or DEFAULT_PROMPT}, ], } diff --git a/tests/test_inference_utils.py b/tests/test_inference_utils.py new file mode 100644 index 0000000..0f82d89 --- /dev/null +++ b/tests/test_inference_utils.py @@ -0,0 +1,58 @@ +import unittest +from unittest.mock import patch + +import numpy as np +import torch + +from moss_transcribe_diarize.inference_utils import ( + DEFAULT_PROMPT, + build_transcription_messages, + load_audio_item, + process_audio_info, +) + + +class InferenceUtilsTest(unittest.TestCase): + @patch("moss_transcribe_diarize.inference_utils.load_audio_item") + def test_process_audio_info_accepts_numpy_waveform(self, load_audio_item): + waveform = np.array([0.1, -0.2, 0.3], dtype=np.float32) + load_audio_item.return_value = waveform + messages = [{"role": "user", "content": [{"type": "audio", "audio": waveform}]}] + + audios = process_audio_info(messages, sampling_rate=16_000) + + self.assertEqual(audios, [waveform]) + load_audio_item.assert_called_once_with(waveform, sampling_rate=16_000) + + @patch("moss_transcribe_diarize.inference_utils.load_audio_item") + def test_process_audio_info_accepts_tensor_without_truth_value_error(self, load_audio_item): + waveform = torch.tensor([0.1, -0.2, 0.3]) + load_audio_item.return_value = waveform.numpy() + messages = [{"role": "user", "content": [{"type": "audio", "audio": waveform}]}] + + process_audio_info(messages, sampling_rate=16_000) + + load_audio_item.assert_called_once_with(waveform, sampling_rate=16_000) + + @patch("moss_transcribe_diarize.inference_utils.load_audio") + def test_load_audio_item_converts_tensor_to_numpy(self, load_audio): + waveform = torch.tensor([0.1, -0.2, 0.3]) + load_audio.return_value = waveform.numpy() + + load_audio_item(waveform, sampling_rate=16_000) + + loaded_audio = load_audio.call_args.args[0] + self.assertIsInstance(loaded_audio, np.ndarray) + np.testing.assert_array_equal(loaded_audio, waveform.numpy()) + + def test_build_transcription_messages_preserves_in_memory_audio(self): + waveform = np.zeros(8, dtype=np.float32) + + messages = build_transcription_messages(waveform, prompt=" ") + + self.assertIs(messages[0]["content"][0]["audio"], waveform) + self.assertEqual(messages[0]["content"][1]["text"], DEFAULT_PROMPT) + + +if __name__ == "__main__": + unittest.main()