Dmitriy Dementiev
0fdeebb256
- Зачем: - необходим устойчивый дефолт модели эмбеддингов и порога на русской речи. - Что: - задокументирован выбор WeSpeaker ResNet34 LM с порогом 0,89. - добавлен возобновляемый скрипт свипа моделей и параметров диаризации. - Проверка: - uvx --cache-dir .uv-cache ruff check scripts/benchmarks/diarization_calibration.py. - uv run --cache-dir .uv-cache pytest: 243 passed, 1 skipped.
356 lines
12 KiB
Python
356 lines
12 KiB
Python
"""Воспроизводимый свип параметров офлайн-диаризации sherpa-onnx."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import itertools
|
|
import json
|
|
import re
|
|
import subprocess
|
|
import time
|
|
import wave
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import sherpa_onnx
|
|
|
|
TURN_RE = re.compile(r"^\*\*\[(\d{2}):(\d{2})(?::(\d{2}))?\] Speaker (\d+):\*\*")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Recording:
|
|
name: str
|
|
path: Path
|
|
start: float
|
|
duration: float
|
|
expected_speakers: int
|
|
reference: Path | None
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--manifest", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--work-dir", type=Path, required=True)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
return parser.parse_args()
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest().upper()
|
|
|
|
|
|
def decode_clip(recording: Recording, work_dir: Path) -> Path:
|
|
output = work_dir / f"{recording.name}.wav"
|
|
if output.exists():
|
|
return output
|
|
|
|
command = [
|
|
"ffmpeg",
|
|
"-hide_banner",
|
|
"-loglevel",
|
|
"error",
|
|
"-y",
|
|
"-ss",
|
|
str(recording.start),
|
|
"-t",
|
|
str(recording.duration),
|
|
"-i",
|
|
str(recording.path),
|
|
"-vn",
|
|
"-ac",
|
|
"1",
|
|
"-ar",
|
|
"16000",
|
|
"-c:a",
|
|
"pcm_s16le",
|
|
str(output),
|
|
]
|
|
subprocess.run(command, check=True)
|
|
return output
|
|
|
|
|
|
def read_wav(path: Path) -> np.ndarray:
|
|
with wave.open(str(path), "rb") as source:
|
|
if source.getnchannels() != 1 or source.getsampwidth() != 2:
|
|
raise ValueError(f"Ожидался mono PCM16 WAV: {path}")
|
|
if source.getframerate() != 16000:
|
|
raise ValueError(f"Ожидалась частота 16 кГц: {path}")
|
|
samples = np.frombuffer(source.readframes(source.getnframes()), np.int16)
|
|
return samples.astype(np.float32) / 32768.0
|
|
|
|
|
|
def timestamp_seconds(match: re.Match[str]) -> float:
|
|
first, second, third = match.group(1), match.group(2), match.group(3)
|
|
if third is None:
|
|
return int(first) * 60 + int(second)
|
|
return int(first) * 3600 + int(second) * 60 + int(third)
|
|
|
|
|
|
def read_reference_turns(recording: Recording) -> list[dict[str, Any]]:
|
|
if recording.reference is None:
|
|
return []
|
|
|
|
starts: list[tuple[float, str]] = []
|
|
for line in recording.reference.read_text(encoding="utf-8").splitlines():
|
|
match = TURN_RE.match(line)
|
|
if match:
|
|
starts.append((timestamp_seconds(match), match.group(4)))
|
|
|
|
clip_end = recording.start + recording.duration
|
|
turns: list[dict[str, Any]] = []
|
|
for index, (start, speaker) in enumerate(starts):
|
|
end = starts[index + 1][0] if index + 1 < len(starts) else clip_end
|
|
overlap_start = max(start, recording.start)
|
|
overlap_end = min(end, clip_end)
|
|
if overlap_end > overlap_start:
|
|
turns.append(
|
|
{
|
|
"speaker": speaker,
|
|
"start": overlap_start - recording.start,
|
|
"end": overlap_end - recording.start,
|
|
}
|
|
)
|
|
return turns
|
|
|
|
|
|
def interval_overlap(left: dict[str, Any], right: dict[str, Any]) -> float:
|
|
return max(0.0, min(left["end"], right["end"]) - max(left["start"], right["start"]))
|
|
|
|
|
|
def best_mapping(
|
|
segments: list[dict[str, Any]],
|
|
reference_turns: list[dict[str, Any]],
|
|
) -> dict[str, Any] | None:
|
|
if not reference_turns or not segments:
|
|
return None
|
|
|
|
predicted = sorted({str(segment["speaker"]) for segment in segments})
|
|
reference = sorted({str(turn["speaker"]) for turn in reference_turns})
|
|
overlap: defaultdict[tuple[str, str], float] = defaultdict(float)
|
|
total = 0.0
|
|
for segment in segments:
|
|
predicted_speaker = str(segment["speaker"])
|
|
for turn in reference_turns:
|
|
value = interval_overlap(segment, turn)
|
|
if value:
|
|
reference_speaker = str(turn["speaker"])
|
|
overlap[(predicted_speaker, reference_speaker)] += value
|
|
total += value
|
|
|
|
best_score = -1.0
|
|
best_pairs: list[tuple[str, str]] = []
|
|
if len(predicted) >= len(reference):
|
|
for candidate in itertools.permutations(predicted, len(reference)):
|
|
pairs = list(zip(candidate, reference, strict=True))
|
|
score = sum(overlap[pair] for pair in pairs)
|
|
if score > best_score:
|
|
best_score, best_pairs = score, pairs
|
|
else:
|
|
for candidate in itertools.permutations(reference, len(predicted)):
|
|
pairs = list(zip(predicted, candidate, strict=True))
|
|
score = sum(overlap[pair] for pair in pairs)
|
|
if score > best_score:
|
|
best_score, best_pairs = score, pairs
|
|
|
|
return {
|
|
"mapped_speaker_purity": best_score / total if total else None,
|
|
"mapped_overlap_seconds": best_score,
|
|
"total_overlap_seconds": total,
|
|
"mapping": {predicted: reference for predicted, reference in best_pairs},
|
|
}
|
|
|
|
|
|
def make_config(
|
|
segmentation_model: Path,
|
|
embedding_model: Path,
|
|
threshold: float,
|
|
num_clusters: int,
|
|
threads: int,
|
|
) -> sherpa_onnx.OfflineSpeakerDiarizationConfig:
|
|
pyannote = sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig(
|
|
model=str(segmentation_model)
|
|
)
|
|
segmentation = sherpa_onnx.OfflineSpeakerSegmentationModelConfig(
|
|
pyannote=pyannote,
|
|
num_threads=threads,
|
|
)
|
|
embedding = sherpa_onnx.SpeakerEmbeddingExtractorConfig(
|
|
model=str(embedding_model),
|
|
num_threads=threads,
|
|
)
|
|
clustering = sherpa_onnx.FastClusteringConfig(
|
|
num_clusters=num_clusters,
|
|
threshold=threshold,
|
|
)
|
|
return sherpa_onnx.OfflineSpeakerDiarizationConfig(
|
|
segmentation=segmentation,
|
|
embedding=embedding,
|
|
clustering=clustering,
|
|
)
|
|
|
|
|
|
def summarize_segments(
|
|
segments: list[dict[str, Any]],
|
|
recording: Recording,
|
|
) -> dict[str, Any]:
|
|
durations: defaultdict[str, float] = defaultdict(float)
|
|
for segment in segments:
|
|
durations[str(segment["speaker"])] += segment["end"] - segment["start"]
|
|
|
|
ordered = sorted(durations.items(), key=lambda item: item[1], reverse=True)
|
|
total = sum(durations.values())
|
|
residual = sum(duration for _, duration in ordered[recording.expected_speakers :])
|
|
substantial_threshold = max(5.0, recording.duration * 0.02)
|
|
return {
|
|
"clusters": len(ordered),
|
|
"substantial_clusters": sum(
|
|
duration >= substantial_threshold for _, duration in ordered
|
|
),
|
|
"substantial_threshold_seconds": substantial_threshold,
|
|
"cluster_durations_seconds": dict(ordered),
|
|
"speaker_time_seconds": total,
|
|
"residual_seconds_after_expected": residual,
|
|
"residual_share_after_expected": residual / total if total else None,
|
|
}
|
|
|
|
|
|
def save_output(path: Path, output: dict[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(
|
|
json.dumps(output, ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
temporary.replace(path)
|
|
|
|
|
|
def manifest_shape(manifest: dict[str, Any]) -> dict[str, Any]:
|
|
"""Отделить параметры эксперимента от машинно-зависимых путей."""
|
|
return {
|
|
"models": [item["name"] for item in manifest["models"]],
|
|
"recordings": [
|
|
{
|
|
key: item[key]
|
|
for key in ("name", "start", "duration", "expected_speakers")
|
|
}
|
|
for item in manifest["recordings"]
|
|
],
|
|
"runs": manifest["runs"],
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
|
|
args.work_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
recordings = [
|
|
Recording(
|
|
name=item["name"],
|
|
path=Path(item["path"]),
|
|
start=float(item["start"]),
|
|
duration=float(item["duration"]),
|
|
expected_speakers=int(item["expected_speakers"]),
|
|
reference=Path(item["reference"]) if item.get("reference") else None,
|
|
)
|
|
for item in manifest["recordings"]
|
|
]
|
|
if args.output.exists():
|
|
output = json.loads(args.output.read_text(encoding="utf-8"))
|
|
if (
|
|
manifest_shape(output["manifest"]) != manifest_shape(manifest)
|
|
or output["threads"] != args.threads
|
|
):
|
|
raise ValueError("Существующий output создан с другим manifest/threads")
|
|
output["manifest"] = manifest
|
|
else:
|
|
output = {
|
|
"manifest": manifest,
|
|
"sherpa_onnx_version": sherpa_onnx.__version__,
|
|
"threads": args.threads,
|
|
"results": [],
|
|
}
|
|
completed = {
|
|
(item["recording"], item["model"], item["run"]) for item in output["results"]
|
|
}
|
|
|
|
segmentation_model = Path(manifest["segmentation_model"])
|
|
for recording in recordings:
|
|
print(f"Декодирование {recording.name}", flush=True)
|
|
wav_path = decode_clip(recording, args.work_dir)
|
|
samples = read_wav(wav_path)
|
|
reference_turns = read_reference_turns(recording)
|
|
source_hash = file_sha256(recording.path)
|
|
|
|
for model in manifest["models"]:
|
|
embedding_model = Path(model["path"])
|
|
for run in manifest["runs"]:
|
|
run_key = (recording.name, model["name"], run["name"])
|
|
if run_key in completed:
|
|
print(f"Пропуск готового прогона: {run_key}", flush=True)
|
|
continue
|
|
num_clusters = run["num_clusters"]
|
|
if num_clusters == "expected":
|
|
num_clusters = recording.expected_speakers
|
|
threshold = float(run["threshold"])
|
|
print(
|
|
f"{recording.name}: {model['name']} / {run['name']}",
|
|
flush=True,
|
|
)
|
|
config = make_config(
|
|
segmentation_model=segmentation_model,
|
|
embedding_model=embedding_model,
|
|
threshold=threshold,
|
|
num_clusters=int(num_clusters),
|
|
threads=args.threads,
|
|
)
|
|
diarizer = sherpa_onnx.OfflineSpeakerDiarization(config)
|
|
started = time.perf_counter()
|
|
result = diarizer.process(samples)
|
|
elapsed = time.perf_counter() - started
|
|
segments = [
|
|
{
|
|
"speaker": int(segment.speaker),
|
|
"start": float(segment.start),
|
|
"end": float(segment.end),
|
|
}
|
|
for segment in result.sort_by_start_time()
|
|
]
|
|
item = {
|
|
"recording": recording.name,
|
|
"source": recording.path.name,
|
|
"source_sha256": source_hash,
|
|
"clip_start": recording.start,
|
|
"clip_duration": recording.duration,
|
|
"expected_speakers": recording.expected_speakers,
|
|
"reference": recording.reference.name
|
|
if recording.reference
|
|
else None,
|
|
"model": model["name"],
|
|
"model_file": embedding_model.name,
|
|
"run": run["name"],
|
|
"threshold": threshold,
|
|
"num_clusters": int(num_clusters),
|
|
"elapsed_seconds": elapsed,
|
|
"rtf": elapsed / recording.duration,
|
|
"summary": summarize_segments(segments, recording),
|
|
"reference_mapping": best_mapping(segments, reference_turns),
|
|
"segments": segments,
|
|
}
|
|
output["results"].append(item)
|
|
completed.add(run_key)
|
|
save_output(args.output, output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|