Files
meeting-lab/src/meeting_lab/transcription/whisper.py
T

291 lines
11 KiB
Python

"""Isolated whisper.cpp wrapper producing Meeting Lab compact transcripts."""
from __future__ import annotations
import json
import os
import platform
import re
import shutil
import subprocess
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Sequence
BACKEND = "whisper.cpp"
RAW_FILENAME = "whisper_raw.json"
TRANSCRIPT_FILENAME = "transcript.json"
TEXT_FILENAME = "transcript.txt"
METADATA_FILENAME = "runtime_metadata.json"
DEFAULT_THREADS = "auto"
class TranscriptionError(RuntimeError):
"""Raised when parameters, Whisper execution, or output are invalid."""
@dataclass(frozen=True)
class TranscriptionResult:
output_dir: Path
raw_output: Path
transcript_json: Path
transcript_text: Path
runtime_metadata: Path
runtime_seconds: float
def _logical_cpu_count() -> int:
"""Return the available logical CPU count as a last-resort fallback."""
if hasattr(os, "sched_getaffinity"):
try:
count = len(os.sched_getaffinity(0))
if count > 0:
return count
except OSError:
pass
return os.cpu_count() or 1
def _linux_physical_core_count() -> int | None:
affinity = None
if hasattr(os, "sched_getaffinity"):
try:
affinity = os.sched_getaffinity(0)
except OSError:
pass
cores: set[tuple[str, str]] = set()
for cpu_dir in Path("/sys/devices/system/cpu").glob("cpu[0-9]*"):
try:
cpu_number = int(cpu_dir.name[3:])
if affinity is not None and cpu_number not in affinity:
continue
topology = cpu_dir / "topology"
package = (topology / "physical_package_id").read_text().strip()
core = (topology / "core_id").read_text().strip()
cores.add((package, core))
except (OSError, ValueError):
continue
return len(cores) or None
def _darwin_physical_core_count() -> int | None:
try:
completed = subprocess.run(
("sysctl", "-n", "hw.physicalcpu"),
check=False,
capture_output=True,
text=True,
)
count = int(completed.stdout.strip())
return count if completed.returncode == 0 and count > 0 else None
except (OSError, ValueError):
return None
def physical_core_count() -> int:
"""Detect physical cores where supported, falling back to available threads."""
system = platform.system()
detected = _linux_physical_core_count() if system == "Linux" else None
if system == "Darwin":
detected = _darwin_physical_core_count()
return detected or _logical_cpu_count()
def resolve_threads(
threads: str | int,
detector: Callable[[], int] = physical_core_count,
) -> int:
if isinstance(threads, bool):
raise TranscriptionError("Threads must be 'auto' or a positive integer.")
if threads == "auto":
count = detector()
else:
try:
count = int(threads)
except (TypeError, ValueError) as exc:
raise TranscriptionError("Threads must be 'auto' or a positive integer.") from exc
if count <= 0:
raise TranscriptionError("Threads must be 'auto' or a positive integer.")
return count
def _vulkan_support(raw: dict[str, Any]) -> bool | None:
system_info = raw.get("systeminfo")
if not isinstance(system_info, str) or "VULKAN" not in system_info.upper():
return None
return re.search(r"VULKAN\s*=\s*1", system_info, re.IGNORECASE) is not None
def _json_object(path: Path) -> dict[str, Any]:
try:
data = json.loads(path.read_text(encoding="utf-8"))
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
raise TranscriptionError(f"Cannot read Whisper JSON output {path}: {exc}") from exc
if not isinstance(data, dict):
raise TranscriptionError("Whisper JSON output must contain a top-level object.")
return data
def compact_transcript(raw: dict[str, Any]) -> dict[str, Any]:
"""Convert whisper.cpp JSON without linguistic cleanup or reordering."""
entries = raw.get("transcription")
if not isinstance(entries, list):
raise TranscriptionError("Whisper JSON output must contain a 'transcription' list.")
segments: list[dict[str, Any]] = []
for index, entry in enumerate(entries):
if not isinstance(entry, dict):
raise TranscriptionError(f"transcription[{index}] must be an object.")
offsets = entry.get("offsets")
if not isinstance(offsets, dict):
raise TranscriptionError(f"transcription[{index}].offsets must be an object.")
start_ms = offsets.get("from")
end_ms = offsets.get("to")
if not isinstance(start_ms, (int, float)) or isinstance(start_ms, bool):
raise TranscriptionError(f"transcription[{index}].offsets.from must be a number.")
if not isinstance(end_ms, (int, float)) or isinstance(end_ms, bool):
raise TranscriptionError(f"transcription[{index}].offsets.to must be a number.")
if end_ms < start_ms:
raise TranscriptionError(
f"transcription[{index}].offsets.to must be greater than or equal to offsets.from."
)
text_value = entry.get("text", "")
if not isinstance(text_value, str):
raise TranscriptionError(f"transcription[{index}].text must be a string.")
text = text_value.strip()
if text:
segments.append(
{
"id": len(segments),
"start": float(start_ms) / 1000.0,
"end": float(end_ms) / 1000.0,
"text": text,
}
)
return {"text": " ".join(item["text"] for item in segments), "segments": segments}
def _timestamp(seconds: float) -> str:
milliseconds = int(round(seconds * 1000))
hours, remainder = divmod(milliseconds, 3_600_000)
minutes, remainder = divmod(remainder, 60_000)
secs, millis = divmod(remainder, 1000)
return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}"
def transcript_text(transcript: dict[str, Any]) -> str:
lines = [
f"[{_timestamp(item['start'])} - {_timestamp(item['end'])}] {item['text']}"
for item in transcript["segments"]
]
return "\n".join(lines) + ("\n" if lines else "")
def _write_json(path: Path, value: dict[str, Any]) -> None:
path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
def transcribe_audio(
audio_path: Path,
model_path: Path,
output_dir: Path,
language: str = "auto",
*,
executable: str = "whisper-cli",
threads: str | int = DEFAULT_THREADS,
thread_detector: Callable[[], int] = physical_core_count,
runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
monotonic: Callable[[], float] = time.monotonic,
now: Callable[[], datetime] = lambda: datetime.now(timezone.utc),
) -> TranscriptionResult:
"""Run one whisper.cpp call and write raw, compact, text, and metadata outputs."""
audio_path = Path(audio_path)
model_path = Path(model_path)
output_dir = Path(output_dir)
if not audio_path.is_file():
raise TranscriptionError(f"Audio file does not exist: {audio_path}")
if not model_path.is_file():
raise TranscriptionError(f"Whisper model does not exist: {model_path}")
if not isinstance(language, str) or not language.strip():
raise TranscriptionError("Language must be a non-empty string.")
if not executable.strip():
raise TranscriptionError("Whisper executable must be a non-empty string.")
if output_dir.exists() and not output_dir.is_dir():
raise TranscriptionError(f"Output directory path is not a directory: {output_dir}")
thread_count = resolve_threads(threads, thread_detector)
output_dir.mkdir(parents=True, exist_ok=True)
raw_output = output_dir / RAW_FILENAME
started_at = now().astimezone(timezone.utc)
started = monotonic()
with tempfile.TemporaryDirectory(prefix=".whisper-", dir=output_dir) as temp_name:
temporary_prefix = Path(temp_name) / "whisper_raw"
command: Sequence[str] = (
executable,
"-m", str(model_path),
"-f", str(audio_path),
"-l", language.strip(),
"-t", str(thread_count),
"-fa",
"-oj",
"-of", str(temporary_prefix),
)
try:
completed = runner(command, check=False, capture_output=True, text=True)
except OSError as exc:
raise TranscriptionError(f"Could not start {BACKEND}: {exc}") from exc
runtime_seconds = monotonic() - started
temporary_raw = temporary_prefix.with_suffix(".json")
if temporary_raw.is_file():
shutil.copyfile(temporary_raw, raw_output)
if completed.returncode != 0:
detail = completed.stderr.strip() or completed.stdout.strip() or "no diagnostic output"
raise TranscriptionError(
f"{BACKEND} failed with exit code {completed.returncode}: {detail}"
)
if not raw_output.is_file():
raise TranscriptionError(f"{BACKEND} completed without producing JSON output.")
raw_data = _json_object(raw_output)
transcript = compact_transcript(raw_data)
transcript_json_path = output_dir / TRANSCRIPT_FILENAME
transcript_text_path = output_dir / TEXT_FILENAME
metadata_path = output_dir / METADATA_FILENAME
_write_json(transcript_json_path, transcript)
transcript_text_path.write_text(transcript_text(transcript), encoding="utf-8")
duration = max((item["end"] for item in transcript["segments"]), default=None)
metadata = {
"input_file": str(audio_path.resolve()),
"model": str(model_path.resolve()),
"backend": BACKEND,
"whisper_executable": executable,
"language": language.strip(),
"threads": thread_count,
"threads_option": str(threads),
"flash_attention": True,
"vulkan_support_detected": _vulkan_support(raw_data),
"duration_seconds": duration,
"runtime_seconds": runtime_seconds,
"timestamp": started_at.isoformat(),
"output_files": {
"whisper_raw": RAW_FILENAME,
"transcript_json": TRANSCRIPT_FILENAME,
"transcript_text": TEXT_FILENAME,
"runtime_metadata": METADATA_FILENAME,
},
}
_write_json(metadata_path, metadata)
return TranscriptionResult(
output_dir=output_dir,
raw_output=raw_output,
transcript_json=transcript_json_path,
transcript_text=transcript_text_path,
runtime_metadata=metadata_path,
runtime_seconds=runtime_seconds,
)