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

chapter6/end-to-end-speech/speech_model.py

@@ -0,0 +1,214 @@
+"""Native local MiniCPM-o 4.5 inference for Experiment 6-5."""
+
+from __future__ import annotations
+
+import hashlib
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+
+MODEL_ID = "openbmb/MiniCPM-o-4_5"
+MODEL_REVISION = "1f761131fa83f5ed3cd6f2f22b225c4501d154fa"
+
+
+@dataclass
+class InferenceResult:
+    mode: str
+    response: str
+    latency_seconds: float
+    transcript: str | None = None
+    stage_latencies: dict[str, float] | None = None
+    output_audio: dict[str, Any] | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def sha256_file(path: str | Path) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def audio_metadata(path: str | Path) -> dict[str, Any]:
+    import soundfile as sf
+
+    path = Path(path)
+    info = sf.info(path)
+    return {
+        "path": str(path),
+        "sha256": sha256_file(path),
+        "sample_rate_hz": info.samplerate,
+        "frames": info.frames,
+        "duration_seconds": round(info.duration, 6),
+        "channels": info.channels,
+        "format": info.format,
+    }
+
+
+class MiniCPMOClient:
+    """One-GPU Transformers client for the official MiniCPM-o 4.5 checkpoint."""
+
+    def __init__(
+        self,
+        model_id: str = MODEL_ID,
+        revision: str = MODEL_REVISION,
+        *,
+        device: str = "cuda",
+        enable_tts: bool = True,
+        local_files_only: bool = False,
+    ) -> None:
+        self.model_id = model_id
+        self.revision = revision
+        self.device = device
+        self.enable_tts = enable_tts
+        self.local_files_only = local_files_only
+        self.model = None
+        self.load_seconds: float | None = None
+
+    def load(self) -> None:
+        import torch
+        from transformers import AutoModel
+
+        if self.device == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError("MiniCPM-o 4.5 local precision run requires an NVIDIA CUDA GPU")
+        started = time.perf_counter()
+        self.model = AutoModel.from_pretrained(
+            self.model_id,
+            revision=self.revision,
+            trust_remote_code=True,
+            attn_implementation="sdpa",
+            torch_dtype=torch.bfloat16,
+            init_vision=False,
+            init_audio=True,
+            init_tts=self.enable_tts,
+            local_files_only=self.local_files_only,
+        )
+        self.model.eval().to(self.device)
+        if self.enable_tts:
+            self.model.init_tts()
+        self.load_seconds = time.perf_counter() - started
+
+    def _require_model(self):
+        if self.model is None:
+            raise RuntimeError("Call load() before inference")
+        return self.model
+
+    @staticmethod
+    def load_audio(path: str | Path):
+        import librosa
+
+        audio, _ = librosa.load(path, sr=16000, mono=True)
+        return audio
+
+    def infer_audio(
+        self,
+        audio_path: str | Path,
+        instruction: str,
+        *,
+        max_new_tokens: int = 256,
+        output_audio_path: str | Path | None = None,
+    ) -> InferenceResult:
+        model = self._require_model()
+        audio = self.load_audio(audio_path)
+        generate_audio = output_audio_path is not None
+        if generate_audio and not self.enable_tts:
+            raise RuntimeError("The client was loaded with enable_tts=False")
+        if output_audio_path is not None:
+            output_audio_path = Path(output_audio_path)
+            output_audio_path.parent.mkdir(parents=True, exist_ok=True)
+        started = time.perf_counter()
+        response = model.chat(
+            msgs=[{"role": "user", "content": [instruction, audio]}],
+            do_sample=False,
+            max_new_tokens=max_new_tokens,
+            use_tts_template=generate_audio,
+            enable_thinking=False,
+            generate_audio=generate_audio,
+            output_audio_path=str(output_audio_path) if output_audio_path else None,
+        )
+        latency = time.perf_counter() - started
+        return InferenceResult(
+            mode="direct-audio-to-speech" if generate_audio else "direct-audio-to-text",
+            response=response,
+            latency_seconds=latency,
+            output_audio=audio_metadata(output_audio_path) if output_audio_path else None,
+        )
+
+    def infer_text(
+        self, prompt: str, *, max_new_tokens: int = 256
+    ) -> InferenceResult:
+        model = self._require_model()
+        started = time.perf_counter()
+        response = model.chat(
+            msgs=[{"role": "user", "content": [prompt]}],
+            do_sample=False,
+            max_new_tokens=max_new_tokens,
+            use_tts_template=False,
+            enable_thinking=False,
+            generate_audio=False,
+        )
+        return InferenceResult(
+            mode="text-only",
+            response=response,
+            latency_seconds=time.perf_counter() - started,
+        )
+
+    def transcribe(self, audio_path: str | Path) -> InferenceResult:
+        return self.infer_audio(
+            audio_path,
+            "Please transcribe only the words spoken in this audio. Do not describe tone, pace, or background sound.",
+            max_new_tokens=256,
+        )
+
+    def self_cascade(
+        self, audio_path: str | Path, instruction: str, *, max_new_tokens: int = 256
+    ) -> InferenceResult:
+        transcription = self.transcribe(audio_path)
+        reasoning = self.infer_text(
+            f"{instruction}\n\nUse only this transcript as evidence:\n{transcription.response}",
+            max_new_tokens=max_new_tokens,
+        )
+        return InferenceResult(
+            mode="self-cascade-audio-to-transcript-to-text",
+            response=reasoning.response,
+            latency_seconds=transcription.latency_seconds + reasoning.latency_seconds,
+            transcript=transcription.response,
+            stage_latencies={
+                "transcription_seconds": transcription.latency_seconds,
+                "reasoning_seconds": reasoning.latency_seconds,
+            },
+        )
+
+    def runtime_metadata(self) -> dict[str, Any]:
+        import torch
+        import transformers
+
+        gpu = None
+        if torch.cuda.is_available():
+            properties = torch.cuda.get_device_properties(0)
+            gpu = {
+                "name": properties.name,
+                "total_memory_gib": round(properties.total_memory / 2**30, 3),
+                "peak_allocated_gib": round(torch.cuda.max_memory_allocated() / 2**30, 3),
+            }
+        return {
+            "model_id": self.model_id,
+            "model_revision": self.revision,
+            "device": self.device,
+            "torch_version": torch.__version__,
+            "transformers_version": transformers.__version__,
+            "cuda_version": torch.version.cuda,
+            "cuda_available": torch.cuda.is_available(),
+            "gpu": gpu,
+            "load_seconds": self.load_seconds,
+            "precision": "bfloat16",
+            "attention": "sdpa",
+            "init_vision": False,
+            "init_audio": True,
+            "init_tts": self.enable_tts,
+        }