ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter6/streaming-speech/qwen2_streaming.py

@@ -0,0 +1,187 @@
+"""Qwen2-Audio growing-prefix inference for Experiment 6-4."""
+
+from __future__ import annotations
+
+import json
+import ast
+import re
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import librosa
+
+
+EVENT_ALIASES = {
+    "noise": "<|noise|>",
+    "background noise": "<|noise|>",
+    "laughter": "<|laughter|>",
+    "laugh": "<|laughter|>",
+    "silence": "<|silence|>",
+    "pause": "<|silence|>",
+    "cough": "<|cough|>",
+}
+
+
+@dataclass
+class PrefixResult:
+    prefix_seconds: float
+    inference_seconds: float
+    transcript: str
+    acoustic_events: list[str]
+    raw_response: str
+
+
+def parse_response(raw: str) -> tuple[str, list[str]]:
+    """Parse the model's requested JSON while retaining raw output for audit."""
+    cleaned = re.sub(r"^```(?:json)?|```$", "", raw.strip()).strip()
+    try:
+        payload = json.loads(cleaned)
+    except json.JSONDecodeError:
+        match = re.search(r"\{.*\}", cleaned, re.DOTALL)
+        payload = None
+        if match:
+            try:
+                payload = json.loads(match.group(0))
+            except json.JSONDecodeError:
+                try:
+                    payload = ast.literal_eval(match.group(0))
+                except (ValueError, SyntaxError):
+                    pass
+        if not isinstance(payload, dict):
+            payload = {"transcript": cleaned, "acoustic_events": []}
+    transcript = str(payload.get("transcript") or "")
+    events: list[str] = []
+    raw_events = payload.get("acoustic_events")
+    if isinstance(raw_events, str):
+        event_list = [raw_events]
+    elif isinstance(raw_events, (list, tuple, set)):
+        event_list = list(raw_events)
+    else:
+        event_list = []
+    for event in event_list:
+        if event is None:
+            continue
+        text = str(event).strip()
+        if not text:
+            continue
+        normalized = EVENT_ALIASES.get(text.lower(), text)
+        if normalized and normalized not in events:
+            events.append(normalized)
+    # Qwen may place an event token next to the transcription rather than in JSON.
+    for token in re.findall(r"<\|[^|]+\|>", raw):
+        if token not in events:
+            events.append(token)
+    return transcript, events
+
+
+class Qwen2AudioRecognizer:
+    """Actual ``Qwen/Qwen2-Audio-7B-Instruct`` inference backend."""
+
+    def __init__(
+        self,
+        model_id: str = "Qwen/Qwen2-Audio-7B-Instruct",
+        device: str = "auto",
+        max_new_tokens: int = 128,
+    ) -> None:
+        self._mlx = device == "mlx" or model_id.startswith("mlx-community/")
+        if self._mlx:
+            from mlx_audio.stt.utils import load_model
+
+            self.device = "mlx"
+            self.model_id = model_id
+            self.max_new_tokens = max_new_tokens
+            self.model = load_model(model_id)
+            self.processor = None
+            return
+        import torch
+        from transformers import AutoProcessor, Qwen2AudioForConditionalGeneration
+
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
+        self.device = device
+        self.model_id = model_id
+        self.max_new_tokens = max_new_tokens
+        dtype = torch.float16 if device in ("cuda", "mps") else torch.float32
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        self.model = Qwen2AudioForConditionalGeneration.from_pretrained(
+            model_id,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+        ).to(device)
+        self.model.eval()
+
+    def transcribe_array(self, audio: Any, sample_rate: int) -> tuple[str, list[str], str]:
+        instruction = (
+            "Transcribe all speech heard so far. Also detect non-verbal acoustic events. "
+            "Return JSON only with keys transcript (string) and acoustic_events (array). "
+            "The acoustic_events array must contain only events actually audible in this clip; "
+            "return an empty array when none are audible. Valid event names are noise, laughter, silence, cough. "
+            "Never copy the list of valid names into the answer. "
+            "Do not treat a pause or background noise as the end of the utterance."
+        )
+        if self._mlx:
+            if sample_rate != 16000:
+                audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
+            result = self.model.generate(
+                audio,
+                prompt=instruction,
+                max_tokens=self.max_new_tokens,
+                temperature=0.0,
+            )
+            raw = result.text.strip()
+            transcript, events = parse_response(raw)
+            return transcript, events, raw
+
+        import torch
+        conversation = [{"role": "user", "content": [
+            {"type": "audio", "audio_url": "prefix.wav"},
+            {"type": "text", "text": instruction},
+        ]}]
+        text = self.processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+        target_sr = self.processor.feature_extractor.sampling_rate
+        if sample_rate != target_sr:
+            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=target_sr)
+        inputs = self.processor(text=text, audios=[audio], return_tensors="pt", padding=True)
+        inputs = {k: v.to(self.device) if hasattr(v, "to") else v for k, v in inputs.items()}
+        with torch.inference_mode():
+            generated = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens, do_sample=False)
+        generated = generated[:, inputs["input_ids"].shape[1]:]
+        raw = self.processor.batch_decode(
+            generated, skip_special_tokens=False, clean_up_tokenization_spaces=False
+        )[0].strip()
+        transcript, events = parse_response(raw)
+        return transcript, events, raw
+
+
+def growing_prefix(
+    recognizer: Qwen2AudioRecognizer,
+    audio_path: str | Path,
+    chunk_seconds: float,
+    *,
+    on_result: Callable[[PrefixResult], None] | None = None,
+) -> list[PrefixResult]:
+    """Re-encode [0:t] for every chunk; this is intentionally not incremental."""
+    audio, sr = librosa.load(str(audio_path), sr=None, mono=True)
+    duration = len(audio) / sr
+    endpoints = []
+    endpoint = chunk_seconds
+    while endpoint < duration:
+        endpoints.append(endpoint)
+        endpoint += chunk_seconds
+    endpoints.append(duration)
+    results = []
+    for endpoint in endpoints:
+        prefix = audio[: max(1, round(endpoint * sr))]
+        started = time.perf_counter()
+        transcript, events, raw = recognizer.transcribe_array(prefix, sr)
+        result = PrefixResult(endpoint, time.perf_counter() - started, transcript, events, raw)
+        results.append(result)
+        if on_result:
+            on_result(result)
+    return results
+
+
+def serialize(results: list[PrefixResult]) -> list[dict[str, Any]]:
+    return [asdict(result) for result in results]