Add initial project structure: pipeline, docs, config, profiles

This commit is contained in:
Кирилл Блинов 2026-05-29 10:16:02 +03:00
parent 4214d689dd
commit 5a5d1fa960
9 changed files with 521 additions and 0 deletions

56
.gitignore vendored Normal file
View File

@ -0,0 +1,56 @@
**/__pycache__/
*.py[cod]
*$py.class
*.so
.Python
env/
venv/
ENV/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
# Модели и кэш
models/
*.bin
*.pt
*.pth
*.onnx
checkpoints/
# Временные и выходные файлы
tmp/
temp/
output/
*.wav
*.mp3
*.m4a
*.ogg
*.flac
*.docx
*.md
*.txt
!.gitkeep
# IDE
.vscode/
.idea/
*.swp
*.swo
*~
# OS
.DS_Store
Thumbs.db

48
config.yaml Normal file
View File

@ -0,0 +1,48 @@
# Конфигурация пайплайна транскрибации совещаний
# Профили оборудования
profiles:
# Текущий: MacBook Air M4, 16GB RAM
mac_m4:
device: cpu # MPS на Mac может быть медленнее/багован; CPU + int8 стабильнее
compute_type: int8
batch_size: 1
model: large-v3
language: ru
diarize: true
# Будущий: GPU с 8GB VRAM
gpu_8gb:
device: cuda
compute_type: float16 # или int8 если не хватает памяти
batch_size: 1 # large-v3 + alignment + diarization в ~8GB
model: large-v3
language: ru
diarize: true
# Универсальный CPU (без GPU)
cpu_best:
device: cpu
compute_type: int8
batch_size: 1
model: large-v3
language: ru
diarize: true
# Активный профиль (можно переопределить через CLI: --profile gpu_8gb)
active_profile: mac_m4
# Настройки диаризации
hf_token: null # HuggingFace токен для pyannote. Установите через env: HF_TOKEN
# Настройки выходного документа
output:
format: docx # docx | md | txt
include_timestamps: true
speaker_label_style: name # name | id | none
paragraph_pause_sec: 2.0 # новый абзац, если пауза > N секунд
# Пути
paths:
output_dir: ./output
temp_dir: ./tmp

1
examples/.gitkeep Normal file
View File

@ -0,0 +1 @@
# Placeholder for example audio files

100
run.py Normal file
View File

@ -0,0 +1,100 @@
"""CLI entrypoint для транскрибации совещаний."""
import argparse
import os
import sys
from pathlib import Path
from src.config import get_profile, load_config
from src.document import build_document
from src.pipeline import run_pipeline
def resolve_device(preferred: str) -> str:
"""Определяет доступное устройство."""
import torch
if preferred == "cuda" and torch.cuda.is_available():
return "cuda"
if preferred == "mps" and torch.backends.mps.is_available():
return "mps"
return "cpu"
def main():
parser = argparse.ArgumentParser(
description="Транскрибация совещаний с диаризацией и таймкодами."
)
parser.add_argument("--input", "-i", required=True, help="Путь к аудиофайлу")
parser.add_argument("--output", "-o", default=None, help="Путь к выходному файлу (docx/md/txt)")
parser.add_argument("--profile", "-p", default=None, help="Профиль конфигурации (mac_m4, gpu_8gb, cpu_best)")
parser.add_argument("--config", "-c", default=None, help="Путь к config.yaml")
parser.add_argument("--device", "-d", default=None, help="Принудительно: cpu/cuda/mps")
parser.add_argument("--model", "-m", default=None, help="Принудительно: tiny/base/small/medium/large-v3")
parser.add_argument("--language", "-l", default=None, help="Язык (ru, en, ...)")
parser.add_argument("--format", "-f", default=None, help="Формат выхода: docx, md, txt")
args = parser.parse_args()
if not Path(args.input).exists():
print(f"Ошибка: файл не найден: {args.input}", file=sys.stderr)
sys.exit(1)
# Загрузка конфига
config = load_config(args.config)
profile = get_profile(config, args.profile)
# Переопределения из CLI
if args.device:
profile["device"] = resolve_device(args.device)
else:
profile["device"] = resolve_device(profile.get("device", "cpu"))
if args.model:
profile["model"] = args.model
if args.language:
profile["language"] = args.language
output_cfg = config.get("output", {})
fmt = args.format or output_cfg.get("format", "docx")
if args.output:
output_path = args.output
else:
stem = Path(args.input).stem
output_dir = Path(config.get("paths", {}).get("output_dir", "./output"))
output_path = str(output_dir / f"{stem}.{fmt}")
# Проверка HF токена
hf_token = os.environ.get("HF_TOKEN") or config.get("hf_token")
if profile.get("diarize", True) and not hf_token:
print(
"Ошибка: для диаризации нужен HuggingFace токен.\n"
"Установите env HF_TOKEN или укажите hf_token в config.yaml",
file=sys.stderr,
)
sys.exit(1)
print(f"Профиль: {args.profile or config.get('active_profile')}")
print(f"Устройство: {profile['device']}")
print(f"Модель: {profile['model']}")
print(f"Язык: {profile['language']}")
print(f"Вход: {args.input}")
print(f"Выход: {output_path}")
print("-" * 40)
# Запуск пайплайна
result = run_pipeline(
audio_path=args.input,
profile_name=args.profile,
config_path=args.config,
output_path=output_path,
)
# Генерация документа
build_document(result["segments"], output_path, config)
print("-" * 40)
print(f"Готово! Сохранено: {output_path}")
if __name__ == "__main__":
main()

0
src/__init__.py Normal file
View File

33
src/config.py Normal file
View File

@ -0,0 +1,33 @@
"""Загрузка и управление конфигурацией."""
import os
from pathlib import Path
from typing import Any, Dict
import yaml
DEFAULT_CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"
def load_config(path: str | Path | None = None) -> Dict[str, Any]:
"""Загружает YAML-конфиг."""
config_path = Path(path) if path else DEFAULT_CONFIG_PATH
with open(config_path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def get_profile(config: Dict[str, Any], profile_name: str | None = None) -> Dict[str, Any]:
"""Возвращает слитый профиль (base + выбранный)."""
active = profile_name or config.get("active_profile", "cpu_best")
profiles = config.get("profiles", {})
if active not in profiles:
raise ValueError(f"Профиль '{active}' не найден. Доступные: {list(profiles.keys())}")
return profiles[active]
def resolve_hf_token(config: Dict[str, Any]) -> str | None:
"""Возвращает HF токен: из конфига или env HF_TOKEN."""
token = config.get("hf_token")
if not token:
token = os.environ.get("HF_TOKEN")
return token

182
src/document.py Normal file
View File

@ -0,0 +1,182 @@
"""Генерация выходных документов (docx, md, txt)."""
import re
from datetime import timedelta
from pathlib import Path
from typing import Any, Dict, List
from docx import Document
from docx.shared import Pt
def format_time(seconds: float) -> str:
"""Форматирует секунды в [HH:MM:SS.mmm]."""
td = timedelta(seconds=seconds)
# td может быть > 1 день, но нам нужно HH:MM:SS
total_seconds = int(td.total_seconds())
hours = total_seconds // 3600
minutes = (total_seconds % 3600) // 60
secs = total_seconds % 60
millis = int((seconds - total_seconds) * 1000)
return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}"
def build_docx(segments: List[Dict[str, Any]], output_path: str, config: Dict[str, Any]) -> str:
"""Собирает .docx с протоколом совещания."""
output_cfg = config.get("output", {})
include_ts = output_cfg.get("include_timestamps", True)
pause_sec = output_cfg.get("paragraph_pause_sec", 2.0)
doc = Document()
doc.add_heading("Протокол совещания", level=0)
# Группируем по спикерам с учётом пауз
current_speaker = None
current_texts: List[str] = []
current_start: float = 0.0
def flush_paragraph():
nonlocal current_speaker, current_texts, current_start
if not current_texts:
return
p = doc.add_paragraph()
if include_ts:
ts = format_time(current_start)
run_ts = p.add_run(f"[{ts}] ")
run_ts.bold = True
run_ts.font.size = Pt(10)
run_spk = p.add_run(f"{current_speaker}:\n")
run_spk.bold = True
run_spk.font.size = Pt(11)
p.add_run(" ".join(current_texts))
current_texts = []
for i, seg in enumerate(segments):
speaker = seg["speaker"]
text = seg["text"]
start = seg["start"]
end = seg["end"]
# Новый абзац, если сменился спикер или пауза большая
prev_end = segments[i - 1]["end"] if i > 0 else start
if speaker != current_speaker or (start - prev_end) > pause_sec:
flush_paragraph()
current_speaker = speaker
current_start = start
if not current_texts:
current_start = start
current_texts.append(text)
flush_paragraph()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
doc.save(output_path)
return output_path
def build_md(segments: List[Dict[str, Any]], output_path: str, config: Dict[str, Any]) -> str:
"""Собирает Markdown."""
output_cfg = config.get("output", {})
include_ts = output_cfg.get("include_timestamps", True)
pause_sec = output_cfg.get("paragraph_pause_sec", 2.0)
lines = ["# Протокол совещания\n", ""]
current_speaker = None
current_texts: List[str] = []
current_start: float = 0.0
def flush():
nonlocal current_speaker, current_texts, current_start
if not current_texts:
return
ts = format_time(current_start)
if include_ts:
lines.append(f"**[{ts}] {current_speaker}:**")
else:
lines.append(f"**{current_speaker}:**")
lines.append(" ".join(current_texts))
lines.append("")
current_texts = []
for i, seg in enumerate(segments):
speaker = seg["speaker"]
text = seg["text"]
start = seg["start"]
prev_end = segments[i - 1]["end"] if i > 0 else start
if speaker != current_speaker or (start - prev_end) > pause_sec:
flush()
current_speaker = speaker
current_start = start
if not current_texts:
current_start = start
current_texts.append(text)
flush()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w", encoding="utf-8") as f:
f.write("\n".join(lines))
return output_path
def build_txt(segments: List[Dict[str, Any]], output_path: str, config: Dict[str, Any]) -> str:
"""Собирает plain text."""
output_cfg = config.get("output", {})
include_ts = output_cfg.get("include_timestamps", True)
pause_sec = output_cfg.get("paragraph_pause_sec", 2.0)
lines: List[str] = ["Протокол совещания\n", ""]
current_speaker = None
current_texts: List[str] = []
current_start: float = 0.0
def flush():
nonlocal current_speaker, current_texts, current_start
if not current_texts:
return
ts = format_time(current_start)
if include_ts:
lines.append(f"[{ts}] {current_speaker}:")
else:
lines.append(f"{current_speaker}:")
lines.append(" ".join(current_texts))
lines.append("")
current_texts = []
for i, seg in enumerate(segments):
speaker = seg["speaker"]
text = seg["text"]
start = seg["start"]
prev_end = segments[i - 1]["end"] if i > 0 else start
if speaker != current_speaker or (start - prev_end) > pause_sec:
flush()
current_speaker = speaker
current_start = start
if not current_texts:
current_start = start
current_texts.append(text)
flush()
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w", encoding="utf-8") as f:
f.write("\n".join(lines))
return output_path
def build_document(segments: List[Dict[str, Any]], output_path: str, config: Dict[str, Any]) -> str:
"""Роутер по формату."""
fmt = Path(output_path).suffix.lower().lstrip(".")
if fmt == "docx":
return build_docx(segments, output_path, config)
elif fmt == "md":
return build_md(segments, output_path, config)
elif fmt == "txt":
return build_txt(segments, output_path, config)
else:
raise ValueError(f"Неподдерживаемый формат: {fmt}. Используйте: docx, md, txt")

101
src/pipeline.py Normal file
View File

@ -0,0 +1,101 @@
"""Основной пайплайн: WhisperX → структурированный результат."""
import os
from pathlib import Path
from typing import Any, Dict, List, Optional
import whisperx
from whisperx.diarize import DiarizationPipeline
from src.config import get_profile, load_config, resolve_hf_token
def run_pipeline(
audio_path: str,
profile_name: Optional[str] = None,
config_path: Optional[str] = None,
output_path: Optional[str] = None,
) -> Dict[str, Any]:
"""
Запускает полный пайплайн транскрибации.
Returns:
Словарь с сегментами (speaker, text, start, end).
"""
config = load_config(config_path)
profile = get_profile(config, profile_name)
device = profile["device"]
compute_type = profile["compute_type"]
batch_size = profile["batch_size"]
model_name = profile["model"]
language = profile["language"]
do_diarize = profile["diarize"]
hf_token = resolve_hf_token(config)
# 1. Загрузка аудио
audio = whisperx.load_audio(audio_path)
# 2. Транскрибация (ASR)
print(f"[Pipeline] Загрузка модели Whisper: {model_name} ({device}, {compute_type})")
model = whisperx.load_model(model_name, device, compute_type=compute_type)
print("[Pipeline] Транскрибация...")
result = model.transcribe(audio, batch_size=batch_size, language=language)
del model
import gc
import torch
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# 3. Alignment (точные таймкоды слов)
print("[Pipeline] Загрузка alignment-модели...")
model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
print("[Pipeline] Выравнивание таймкодов...")
result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
del model_a
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# 4. Диаризация (кто говорил)
if do_diarize:
if not hf_token:
raise RuntimeError("HF_TOKEN не задан. Укажите в config.yaml или env HF_TOKEN.")
print("[Pipeline] Загрузка модели диаризации...")
diarize_model = DiarizationPipeline(token=hf_token, device=device)
print("[Pipeline] Диаризация...")
diarize_segments = diarize_model(audio)
result = whisperx.assign_word_speakers(diarize_segments, result)
del diarize_model
gc.collect()
if device == "cuda":
torch.cuda.empty_cache()
# 5. Форматирование результата
segments = format_segments(result)
return {
"segments": segments,
"language": result.get("language", language),
"output_path": output_path,
}
def format_segments(result: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Преобразует результат whisperx в плоский список сегментов."""
segments = []
for seg in result.get("segments", []):
speaker = seg.get("speaker", "UNKNOWN")
text = seg.get("text", "").strip()
start = seg.get("start", 0.0)
end = seg.get("end", 0.0)
if text:
segments.append({
"speaker": speaker,
"text": text,
"start": start,
"end": end,
})
return segments

0
tests/__init__.py Normal file
View File