transcription/run.py

101 lines
3.7 KiB
Python
Raw Normal View History

"""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()